use serde::de::DeserializeOwned;
use serde::{Deserialize, Deserializer, Serialize};
use sha2::{Digest, Sha256};
use std::collections::{BTreeMap, BTreeSet};
use super::{
checked_elements, physical_component_ids, validate_physical_layout_budget, AttributeId,
AxisWeightComponent, BlockQuantizationSpec, CanonicalRational, CompositeWeightPart,
ContractVersion, ElementType, ExternalModelMetadataId, ModelFamilyId, NodeId, OperationId,
PhysicalStorageLayout, PhysicalWeightComponentBinding, PhysicalWeightLayout,
PhysicalWeightPadding, ProgramValueId, QuantizationGrouping, QuantizationPacking,
QuantizationSpec, ResolvedTensorLayout, ResolvedWeightBinding, ResolvedWeightComponentLayout,
ResolvedWeightLogicalValidation, SemanticValue, StateId, StateInitialization, TokenizerId,
VNextError, WeightComponentRole, WeightEncoding, WeightFormatId, WeightId, WeightLayoutId,
MAX_PHYSICAL_WEIGHT_LAYOUT_DEPTH, MAX_PHYSICAL_WEIGHT_LAYOUT_NODES,
};
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct WeightComponentSpec {
pub id: WeightId,
pub role: WeightComponentRole,
pub external_names: Vec<String>,
pub dimensions: Vec<u64>,
pub encoding: WeightEncoding,
pub required: bool,
}
impl WeightComponentSpec {
pub fn physical_bytes(&self) -> Result<u64, VNextError> {
self.encoding.physical_bytes(&self.dimensions, &self.id)
}
pub fn dense_element_type(&self) -> Option<ElementType> {
self.encoding.dense_element_type()
}
pub fn physical_element_type(&self) -> ElementType {
self.dense_element_type().unwrap_or(ElementType::U8)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct PhysicalWeightComponentRef {
pub component_id: WeightId,
pub physical_dimensions: Vec<u64>,
pub resource_bytes: u64,
pub element_type: ElementType,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct WeightTensorSpec {
pub id: WeightId,
pub dimensions: Vec<u64>,
pub logical_element_type: ElementType,
pub physical_layout: PhysicalWeightLayout,
pub required: bool,
}
impl WeightTensorSpec {
pub fn logical_elements(&self) -> Result<u64, VNextError> {
checked_elements(&self.dimensions).ok_or_else(|| VNextError::InvalidExecutionPlan {
reason: format!("logical weight `{}` element count overflows u64", self.id),
})
}
pub fn logical_bytes(&self) -> Result<u64, VNextError> {
self.logical_elements()?
.checked_mul(self.logical_element_type.size_bytes())
.ok_or_else(|| VNextError::InvalidExecutionPlan {
reason: format!("logical weight `{}` byte size overflows u64", self.id),
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct WeightSchema {
pub format_id: WeightFormatId,
pub layout_id: WeightLayoutId,
pub version: ContractVersion,
pub components: Vec<WeightComponentSpec>,
pub tensors: Vec<WeightTensorSpec>,
}
impl WeightSchema {
pub(crate) fn normalize(&mut self) {
self.components
.sort_by(|left, right| left.id.cmp(&right.id));
for tensor in &mut self.tensors {
tensor.physical_layout.normalize();
}
self.tensors.sort_by(|left, right| left.id.cmp(&right.id));
}
pub fn validate(&self, family_id: &ModelFamilyId) -> Result<(), VNextError> {
if self.version.major == 0 || self.components.is_empty() || self.tensors.is_empty() {
return Err(VNextError::UnknownWeightLayout {
family_id: family_id.to_string(),
layout_id: self.layout_id.to_string(),
});
}
for tensor in &self.tensors {
validate_physical_layout_budget(&tensor.physical_layout).map_err(|reason| {
VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: format!("weight_schema.tensors.{}.physical_layout", tensor.id),
reason,
}
})?;
}
let mut component_ids = BTreeSet::new();
let mut names = BTreeSet::new();
let mut components = BTreeMap::new();
let mut quantization_abis = BTreeMap::new();
for component in &self.components {
if !component_ids.insert(component.id.clone())
|| component.external_names.is_empty()
|| component.dimensions.is_empty()
|| component
.external_names
.iter()
.any(|name| name.trim().is_empty() || !names.insert(name.clone()))
|| component.dimensions.iter().any(|extent| *extent == 0)
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "weight_schema.components".to_owned(),
reason: "component identities, names, and dimensions must be valid and unique"
.to_owned(),
});
}
if let WeightEncoding::Quantized(quantization) = &component.encoding {
quantization.validate()?;
}
if let WeightEncoding::BlockQuantized(quantization) = &component.encoding {
quantization.validate()?;
}
let quantization_format = match &component.encoding {
WeightEncoding::Quantized(spec) => Some(&spec.format_id),
WeightEncoding::BlockQuantized(spec) => Some(&spec.format_id),
WeightEncoding::Dense { .. } | WeightEncoding::DenseAffine { .. } => None,
};
if let Some(format_id) = quantization_format {
if let Some(existing) = quantization_abis.get(format_id) {
if existing != &component.encoding {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "weight_schema.components.encoding".to_owned(),
reason: format!(
"quantization format `{format_id}` maps to conflicting physical ABIs"
),
});
}
} else {
quantization_abis.insert(format_id.clone(), component.encoding.clone());
}
}
if let WeightEncoding::DenseAffine { element_type, .. } = &component.encoding {
if component.role != WeightComponentRole::Values
|| !matches!(
element_type,
ElementType::F16 | ElementType::Bf16 | ElementType::F32
)
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "weight_schema.components.encoding".to_owned(),
reason: format!(
"affine dense component `{}` must be a floating-point value component",
component.id
),
});
}
}
component
.physical_bytes()
.map_err(|error| VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "weight_schema.components.dimensions".to_owned(),
reason: error.to_string(),
})?;
let role_encoding_valid = match component.role {
WeightComponentRole::Scales => matches!(
component.encoding,
WeightEncoding::Dense {
element_type: ElementType::U8
| ElementType::F16
| ElementType::Bf16
| ElementType::F32
}
),
WeightComponentRole::ZeroPoints
| WeightComponentRole::Indices
| WeightComponentRole::Permutation => matches!(
component.encoding,
WeightEncoding::Dense {
element_type: ElementType::U8
| ElementType::U32
| ElementType::I8
| ElementType::I32
}
),
WeightComponentRole::PackedValues => {
matches!(
component.encoding,
WeightEncoding::Quantized(_) | WeightEncoding::BlockQuantized(_)
)
}
_ => true,
};
if !role_encoding_valid {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "weight_schema.components.encoding".to_owned(),
reason: format!(
"component `{}` encoding is incompatible with its structural role",
component.id
),
});
}
components.insert(component.id.clone(), component);
}
let mut tensor_ids = BTreeSet::new();
let mut referenced_components = BTreeSet::new();
for tensor in &self.tensors {
if !tensor_ids.insert(tensor.id.clone())
|| tensor.dimensions.is_empty()
|| tensor.dimensions.iter().any(|extent| *extent == 0)
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "weight_schema.tensors".to_owned(),
reason: "logical weight identities and dimensions must be valid and unique"
.to_owned(),
});
}
self.validate_physical_layout(
family_id,
tensor,
&components,
&mut referenced_components,
)?;
}
if let Some(component) = self
.components
.iter()
.find(|component| component.required && !referenced_components.contains(&component.id))
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "weight_schema.components".to_owned(),
reason: format!(
"required component `{}` is not referenced by a logical weight",
component.id
),
});
}
Ok(())
}
pub fn fingerprint(&self) -> Result<String, VNextError> {
let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
context: "fingerprint weight schema",
message: error.to_string(),
})?;
Ok(format!("{:x}", Sha256::digest(bytes)))
}
pub fn quantization_formats(&self) -> BTreeSet<super::QuantizationFormatId> {
self.components
.iter()
.filter_map(|component| match &component.encoding {
WeightEncoding::Quantized(spec) => Some(spec.format_id.clone()),
WeightEncoding::BlockQuantized(spec) => Some(spec.format_id.clone()),
WeightEncoding::Dense { .. } | WeightEncoding::DenseAffine { .. } => None,
})
.collect()
}
fn validate_physical_layout(
&self,
family_id: &ModelFamilyId,
tensor: &WeightTensorSpec,
components: &BTreeMap<WeightId, &WeightComponentSpec>,
referenced: &mut BTreeSet<WeightId>,
) -> Result<(), VNextError> {
let mut validator = PhysicalLayoutValidator {
family_id,
tensor_id: &tensor.id,
components,
referenced,
visited_nodes: 0,
};
validator.validate_layout(
&tensor.physical_layout,
&tensor.dimensions,
tensor.logical_element_type,
1,
)
}
pub fn tensor(&self, weight_id: &WeightId) -> Option<&WeightTensorSpec> {
self.tensors.iter().find(|tensor| &tensor.id == weight_id)
}
pub fn physical_component_refs(
&self,
weight_id: &WeightId,
) -> Result<Vec<&WeightComponentSpec>, VNextError> {
let tensor = self
.tensor(weight_id)
.ok_or_else(|| VNextError::InvalidExecutionPlan {
reason: format!("unknown logical weight `{weight_id}`"),
})?;
let required = physical_component_ids(&tensor.physical_layout).map_err(|reason| {
VNextError::InvalidExecutionPlan {
reason: format!(
"logical weight `{weight_id}` has invalid physical layout: {reason}"
),
}
})?;
let result = self
.components
.iter()
.filter(|component| required.contains(&component.id))
.collect::<Vec<_>>();
if result.len() != required.len() {
return Err(VNextError::InvalidExecutionPlan {
reason: format!("logical weight `{weight_id}` references an unknown component"),
});
}
Ok(result)
}
pub fn physical_bytes(&self, weight_id: &WeightId) -> Result<u64, VNextError> {
self.physical_component_refs(weight_id)?
.into_iter()
.try_fold(0_u64, |total, component| {
total
.checked_add(component.physical_bytes()?)
.ok_or_else(|| VNextError::InvalidExecutionPlan {
reason: format!("logical weight `{weight_id}` physical bytes overflow u64"),
})
})
}
pub fn physical_resource_requirements(
&self,
weight_id: &WeightId,
) -> Result<Vec<PhysicalWeightComponentRef>, VNextError> {
self.physical_component_refs(weight_id)?
.into_iter()
.map(|component| {
Ok(PhysicalWeightComponentRef {
component_id: component.id.clone(),
physical_dimensions: component.dimensions.clone(),
resource_bytes: component.physical_bytes()?,
element_type: component.physical_element_type(),
})
})
.collect()
}
}
impl ResolvedWeightBinding {
pub fn from_schema(schema: &WeightSchema, weight_id: &WeightId) -> Result<Self, VNextError> {
let tensor = schema
.tensor(weight_id)
.ok_or_else(|| VNextError::InvalidExecutionPlan {
reason: format!("unknown logical weight `{weight_id}`"),
})?;
let mut components = schema
.physical_component_refs(weight_id)?
.into_iter()
.map(|component| {
ResolvedWeightComponentLayout::from_parts(
component.id.clone(),
component.role,
component.dimensions.clone(),
component.encoding.clone(),
)
})
.collect::<Vec<_>>();
components.sort_by(|left, right| left.component_id().cmp(right.component_id()));
let binding = Self::from_parts(
weight_id.clone(),
schema.format_id.clone(),
schema.layout_id.clone(),
schema.version,
tensor.physical_layout.clone(),
components,
)?;
binding.validate_logical(&tensor.dimensions, tensor.logical_element_type)?;
Ok(binding)
}
pub fn validate_logical(
&self,
logical_dimensions: &[u64],
logical_element_type: ElementType,
) -> Result<(), VNextError> {
ResolvedWeightLogicalValidation::validate_logical_contract(
self,
logical_dimensions,
logical_element_type,
)
}
}
impl ResolvedWeightLogicalValidation for ResolvedWeightBinding {
fn validate_logical_contract(
&self,
logical_dimensions: &[u64],
logical_element_type: ElementType,
) -> Result<(), VNextError> {
self.validate_structure()?;
let schema = WeightSchema {
format_id: self.schema_format_id().clone(),
layout_id: self.layout_id().clone(),
version: self.schema_version(),
components: self
.components()
.iter()
.map(|component| WeightComponentSpec {
id: component.component_id().clone(),
role: component.role(),
external_names: vec![format!("resolved.{}", component.component_id())],
dimensions: component.physical_dimensions().to_vec(),
encoding: component.encoding().clone(),
required: true,
})
.collect(),
tensors: vec![WeightTensorSpec {
id: self.weight_id().clone(),
dimensions: logical_dimensions.to_vec(),
logical_element_type,
physical_layout: self.physical_layout().clone(),
required: true,
}],
};
schema.validate(&ModelFamilyId::new("family.resolved-weight-binding")?)
}
}
struct PhysicalLayoutValidator<'schema, 'references> {
family_id: &'schema ModelFamilyId,
tensor_id: &'schema WeightId,
components: &'schema BTreeMap<WeightId, &'schema WeightComponentSpec>,
referenced: &'references mut BTreeSet<WeightId>,
visited_nodes: usize,
}
impl<'schema, 'references> PhysicalLayoutValidator<'schema, 'references> {
fn invalid(&self, reason: impl Into<String>) -> VNextError {
VNextError::InvalidModelConfig {
family_id: self.family_id.to_string(),
field: format!("weight_schema.tensors.{}.physical_layout", self.tensor_id),
reason: reason.into(),
}
}
fn visit_node(&mut self, depth: usize) -> Result<(), VNextError> {
if depth > MAX_PHYSICAL_WEIGHT_LAYOUT_DEPTH {
return Err(self.invalid(format!(
"physical layout depth exceeds {MAX_PHYSICAL_WEIGHT_LAYOUT_DEPTH}"
)));
}
self.visited_nodes = self
.visited_nodes
.checked_add(1)
.ok_or_else(|| self.invalid("physical layout node count overflows usize"))?;
if self.visited_nodes > MAX_PHYSICAL_WEIGHT_LAYOUT_NODES {
return Err(self.invalid(format!(
"physical layout node count exceeds {MAX_PHYSICAL_WEIGHT_LAYOUT_NODES}"
)));
}
Ok(())
}
fn component(
&self,
component_id: &WeightId,
) -> Result<&'schema WeightComponentSpec, VNextError> {
self.components
.get(component_id)
.copied()
.ok_or_else(|| self.invalid(format!("unknown component `{component_id}`")))
}
fn bind_component(
&mut self,
binding: &PhysicalWeightComponentBinding,
semantic_dimensions: &[u64],
role: WeightComponentRole,
depth: usize,
) -> Result<&'schema WeightComponentSpec, VNextError> {
self.visit_node(depth)?;
let component = self.component(&binding.component_id)?;
if component.role != role {
return Err(self.invalid(format!(
"component `{}` has role {:?}, expected {:?}",
component.id, component.role, role
)));
}
self.validate_storage(component, semantic_dimensions, &binding.storage)?;
if !self.referenced.insert(component.id.clone()) {
return Err(self.invalid(format!(
"component `{}` is referenced more than once in the physical layout tree",
component.id
)));
}
Ok(component)
}
fn validate_storage(
&self,
component: &WeightComponentSpec,
semantic_dimensions: &[u64],
storage: &PhysicalStorageLayout,
) -> Result<(), VNextError> {
if semantic_dimensions.is_empty()
|| semantic_dimensions.iter().any(|extent| *extent == 0)
|| checked_elements(semantic_dimensions).is_none()
{
return Err(self.invalid(format!(
"component `{}` has an invalid or overflowing semantic shape",
component.id
)));
}
let raw_elements = checked_elements(&component.dimensions).ok_or_else(|| {
self.invalid(format!(
"component `{}` raw storage shape overflows u64",
component.id
))
})?;
match storage {
PhysicalStorageLayout::Contiguous { padding } => {
let padded = self.resolve_padding(semantic_dimensions, padding)?;
if component.dimensions != padded {
return Err(self.invalid(format!(
"component `{}` contiguous shape {:?} differs from its explicit physical shape {:?}",
component.id, padded, component.dimensions
)));
}
}
PhysicalStorageLayout::Strided {
strides_in_elements,
padding,
} => {
let padded = self.resolve_padding(semantic_dimensions, padding)?;
let span = self.checked_strided_span(&padded, strides_in_elements, 1)?;
if span != raw_elements {
return Err(self.invalid(format!(
"component `{}` strided span {span} differs from its raw storage element count {raw_elements}",
component.id
)));
}
}
PhysicalStorageLayout::Tiled {
tile_shape,
axis_order,
tile_strides_in_elements,
padding,
} => {
let rank = semantic_dimensions.len();
if tile_shape.len() != rank
|| tile_shape.iter().any(|extent| *extent == 0)
|| !is_axis_permutation(axis_order, rank)
|| tile_strides_in_elements.len() != rank
{
return Err(self.invalid(format!(
"component `{}` tile shape, axis order, or strides do not match rank",
component.id
)));
}
let padded = self.resolve_padding(semantic_dimensions, padding)?;
let minimal_padded = semantic_dimensions
.iter()
.zip(tile_shape)
.map(|(extent, tile)| checked_round_up(*extent, *tile))
.collect::<Option<Vec<_>>>()
.ok_or_else(|| {
self.invalid(format!(
"component `{}` tile padding overflows u64",
component.id
))
})?;
match padding {
PhysicalWeightPadding::Exact if minimal_padded != semantic_dimensions => {
return Err(self.invalid(format!(
"component `{}` needs tile padding but declares exact storage",
component.id
)));
}
PhysicalWeightPadding::ZeroFill { .. } if padded != minimal_padded => {
return Err(self.invalid(format!(
"component `{}` tiled zero-fill shape is not the unique minimal padded shape",
component.id
)));
}
_ => {}
}
let semantic_grid = padded
.iter()
.zip(tile_shape)
.map(|(extent, tile)| extent / tile)
.collect::<Vec<_>>();
let physical_grid = axis_order
.iter()
.map(|axis| semantic_grid[*axis as usize])
.collect::<Vec<_>>();
let tile_elements = checked_elements(tile_shape).ok_or_else(|| {
self.invalid(format!(
"component `{}` tile size overflows u64",
component.id
))
})?;
let span = self.checked_strided_span(
&physical_grid,
tile_strides_in_elements,
tile_elements,
)?;
if span != raw_elements {
return Err(self.invalid(format!(
"component `{}` tiled span {span} differs from its raw storage element count {raw_elements}",
component.id
)));
}
}
}
Ok(())
}
fn resolve_padding(
&self,
semantic_dimensions: &[u64],
padding: &PhysicalWeightPadding,
) -> Result<Vec<u64>, VNextError> {
match padding {
PhysicalWeightPadding::Exact => Ok(semantic_dimensions.to_vec()),
PhysicalWeightPadding::ZeroFill { padded_dimensions } => {
if padded_dimensions.len() != semantic_dimensions.len()
|| padded_dimensions.iter().any(|extent| *extent == 0)
|| padded_dimensions
.iter()
.zip(semantic_dimensions)
.any(|(padded, semantic)| padded < semantic)
|| padded_dimensions == semantic_dimensions
|| checked_elements(padded_dimensions).is_none()
{
return Err(self.invalid(
"zero-fill padding must explicitly enlarge a valid shape without shrinking any axis",
));
}
Ok(padded_dimensions.clone())
}
}
}
fn checked_strided_span(
&self,
dimensions: &[u64],
strides: &[u64],
base_span: u64,
) -> Result<u64, VNextError> {
if dimensions.is_empty()
|| dimensions.len() != strides.len()
|| dimensions.iter().any(|extent| *extent == 0)
|| strides.iter().any(|stride| *stride == 0)
|| base_span == 0
{
return Err(self.invalid("strided storage dimensions and strides are invalid"));
}
let mut axes = dimensions
.iter()
.copied()
.zip(strides.iter().copied())
.filter(|(extent, _)| *extent > 1)
.collect::<Vec<_>>();
axes.sort_by_key(|(_, stride)| *stride);
let mut span = base_span;
for (extent, stride) in axes {
if stride < span {
return Err(
self.invalid("strided storage aliases coordinates or overlaps physical tiles")
);
}
span = extent
.checked_sub(1)
.and_then(|count| count.checked_mul(stride))
.and_then(|addition| span.checked_add(addition))
.ok_or_else(|| self.invalid("strided storage span overflows u64"))?;
}
Ok(span)
}
fn grouped_dimensions(
&self,
semantic_dimensions: &[u64],
padding: &PhysicalWeightPadding,
group_axis: usize,
group_size: u64,
) -> Result<Vec<u64>, VNextError> {
let axis_extent = semantic_dimensions[group_axis];
let minimal_axis = checked_round_up(axis_extent, group_size)
.ok_or_else(|| self.invalid("quantization group padding overflows u64"))?;
match padding {
PhysicalWeightPadding::Exact => {
if minimal_axis != axis_extent {
return Err(self.invalid(
"quantization groups require padding but exact storage was declared",
));
}
Ok(semantic_dimensions.to_vec())
}
PhysicalWeightPadding::ZeroFill { padded_dimensions } => {
if minimal_axis == axis_extent
|| padded_dimensions.len() != semantic_dimensions.len()
|| padded_dimensions.iter().enumerate().any(|(axis, extent)| {
if axis == group_axis {
*extent != minimal_axis
} else {
*extent != semantic_dimensions[axis]
}
})
{
return Err(self.invalid(
"quantization zero-fill must pad only the group axis to its unique minimal extent",
));
}
checked_elements(padded_dimensions)
.is_some()
.then(|| padded_dimensions.clone())
.ok_or_else(|| self.invalid("quantization padded shape overflows u64"))
}
}
}
fn validate_dense_values(
&mut self,
binding: &PhysicalWeightComponentBinding,
semantic_dimensions: &[u64],
logical_element_type: ElementType,
depth: usize,
) -> Result<(), VNextError> {
let component = self.bind_component(
binding,
semantic_dimensions,
WeightComponentRole::Values,
depth,
)?;
if component.dense_element_type() != Some(logical_element_type) {
return Err(self.invalid(format!(
"values component `{}` dtype differs from the logical tensor",
component.id
)));
}
Ok(())
}
fn validate_axis_component(
&mut self,
axis_component: &AxisWeightComponent,
semantic_dimensions: &[u64],
expected_axis: usize,
role: WeightComponentRole,
allow_narrow_integer: bool,
depth: usize,
) -> Result<(), VNextError> {
if axis_component.axis as usize != expected_axis {
return Err(self.invalid(format!(
"axis component `{}` targets axis {}, expected {expected_axis}",
axis_component.component.component_id, axis_component.axis
)));
}
let axis_shape = [semantic_dimensions[expected_axis]];
let component = self.bind_component(&axis_component.component, &axis_shape, role, depth)?;
let integer_type_valid = component.dense_element_type().is_some_and(|element_type| {
if allow_narrow_integer {
matches!(
element_type,
ElementType::U8 | ElementType::U32 | ElementType::I8 | ElementType::I32
)
} else {
matches!(element_type, ElementType::U32 | ElementType::I32)
}
});
if !integer_type_valid {
return Err(self.invalid(format!(
"axis component `{}` must use an integer encoding valid for {:?}",
component.id, role
)));
}
Ok(())
}
fn validate_layout(
&mut self,
layout: &PhysicalWeightLayout,
semantic_dimensions: &[u64],
logical_element_type: ElementType,
depth: usize,
) -> Result<(), VNextError> {
self.visit_node(depth)?;
if semantic_dimensions.is_empty()
|| semantic_dimensions.iter().any(|extent| *extent == 0)
|| checked_elements(semantic_dimensions).is_none()
{
return Err(self.invalid("logical layout shape is empty, zero, or overflowing"));
}
match layout {
PhysicalWeightLayout::Dense { component_id } => {
let binding =
PhysicalWeightComponentBinding::exact_contiguous(component_id.clone());
self.validate_dense_values(
&binding,
semantic_dimensions,
logical_element_type,
depth,
)?;
}
PhysicalWeightLayout::Stored { component } => {
self.validate_dense_values(
component,
semantic_dimensions,
logical_element_type,
depth,
)?;
}
PhysicalWeightLayout::Composite { parts } => {
if parts.is_empty() {
return Err(self.invalid("composite layout has no parts"));
}
let rank = semantic_dimensions.len();
let mut covered_elements = 0_u64;
for (index, part) in parts.iter().enumerate() {
if part.logical_offsets.len() != rank
|| part.extents.len() != rank
|| part.extents.iter().any(|extent| *extent == 0)
|| part
.logical_offsets
.iter()
.zip(&part.extents)
.zip(semantic_dimensions)
.any(|((offset, extent), logical)| {
offset.checked_add(*extent).is_none_or(|end| end > *logical)
})
{
return Err(self.invalid(format!(
"composite part {index} has invalid semantic offsets or extents"
)));
}
for previous in &parts[..index] {
let overlaps = part
.logical_offsets
.iter()
.zip(&part.extents)
.zip(previous.logical_offsets.iter().zip(&previous.extents))
.all(|((offset, extent), (other_offset, other_extent))| {
offset.checked_add(*extent).is_some_and(|end| {
other_offset.checked_add(*other_extent).is_some_and(
|other_end| *offset < other_end && *other_offset < end,
)
})
});
if overlaps {
return Err(self.invalid("composite semantic placements overlap"));
}
}
let part_elements = checked_elements(&part.extents)
.ok_or_else(|| self.invalid("composite part size overflows u64"))?;
covered_elements = covered_elements
.checked_add(part_elements)
.ok_or_else(|| self.invalid("composite coverage overflows u64"))?;
self.validate_layout(
&part.layout,
&part.extents,
logical_element_type,
depth + 1,
)?;
}
if covered_elements != checked_elements(semantic_dimensions).unwrap() {
return Err(self.invalid(
"composite semantic placements do not cover the logical tensor exactly",
));
}
}
PhysicalWeightLayout::Quantized {
packed_values,
packed_dimensions,
scales,
zero_points,
zero_point_packed_dimensions,
axis_indices,
permutation,
codebook,
group_axis,
group_padding,
} => {
if !matches!(
logical_element_type,
ElementType::F16 | ElementType::Bf16 | ElementType::F32
) {
return Err(
self.invalid("quantized logical weight dtype must be floating point")
);
}
let axis = *group_axis as usize;
if axis >= semantic_dimensions.len() {
return Err(self.invalid("quantization group axis is out of range"));
}
let quantization = {
let component = self.component(&packed_values.component_id)?;
let WeightEncoding::Quantized(spec) = &component.encoding else {
return Err(self.invalid(
"packed-values component does not carry a quantization spec",
));
};
spec.clone()
};
let group_size = quantization
.grouping
.resolved_size(semantic_dimensions[axis]);
if group_size == 0 {
return Err(self.invalid(
"two-dimensional block grouping requires a block-grid quantized layout",
));
}
let grouped_dimensions =
self.grouped_dimensions(semantic_dimensions, group_padding, axis, group_size)?;
let packed_bytes = checked_elements(&grouped_dimensions)
.and_then(|elements| {
elements.checked_mul(u64::from(quantization.bits_per_weight))
})
.and_then(|bits| bits.checked_add(7))
.map(|bits| bits / 8)
.ok_or_else(|| self.invalid("packed-values size overflows u64"))?;
if checked_elements(packed_dimensions) != Some(packed_bytes) {
return Err(self.invalid(format!(
"packed-values semantic shape contains {} storage bytes, expected {packed_bytes}",
checked_elements(packed_dimensions)
.map_or_else(|| "an overflowing number of".to_owned(), |value| value.to_string())
)));
}
let packed = self.bind_component(
packed_values,
packed_dimensions,
WeightComponentRole::PackedValues,
depth,
)?;
if packed.encoding != WeightEncoding::Quantized(quantization.clone()) {
return Err(self.invalid(
"packed-values encoding changed while validating the quantized tree",
));
}
let mut group_shape = grouped_dimensions;
group_shape[axis] /= group_size;
let scales_component =
self.bind_component(scales, &group_shape, WeightComponentRole::Scales, depth)?;
if scales_component.dense_element_type() != Some(quantization.scale_type) {
return Err(
self.invalid("scale component dtype differs from the quantization spec")
);
}
match (
quantization.zero_point_type,
zero_points,
zero_point_packed_dimensions,
) {
(Some(expected_type), Some(binding), Some(packed_dimensions)) => {
let expected_bytes = checked_elements(&group_shape)
.and_then(|elements| {
elements.checked_mul(u64::from(quantization.bits_per_weight))
})
.and_then(|bits| bits.checked_add(7))
.map(|bits| bits / 8)
.ok_or_else(|| self.invalid("packed zero-point size overflows u64"))?;
let component = self.bind_component(
binding,
packed_dimensions,
WeightComponentRole::ZeroPoints,
depth,
)?;
if component.dense_element_type() != Some(expected_type)
|| component.physical_bytes()? != expected_bytes
{
return Err(self.invalid(
"packed zero-point component differs from its quantization contract",
));
}
}
(Some(expected_type), Some(binding), None) => {
let component = self.bind_component(
binding,
&group_shape,
WeightComponentRole::ZeroPoints,
depth,
)?;
if component.dense_element_type() != Some(expected_type) {
return Err(self.invalid(
"zero-point component dtype differs from the quantization spec",
));
}
}
(None, None, None) => {}
_ => {
return Err(self.invalid(
"zero-point component presence differs from the quantization spec",
));
}
}
if let Some(axis_indices) = axis_indices {
self.validate_axis_component(
axis_indices,
semantic_dimensions,
axis,
WeightComponentRole::Indices,
true,
depth,
)?;
}
if let Some(permutation) = permutation {
self.validate_axis_component(
permutation,
semantic_dimensions,
axis,
WeightComponentRole::Permutation,
false,
depth,
)?;
}
if let Some(codebook) = codebook {
let entries = 1_u64
.checked_shl(u32::from(quantization.bits_per_weight))
.ok_or_else(|| self.invalid("codebook size overflows u64"))?;
let component = self.bind_component(
codebook,
&[entries],
WeightComponentRole::Codebook,
depth,
)?;
if component.dense_element_type() != Some(logical_element_type) {
return Err(
self.invalid("codebook dtype differs from the logical tensor dtype")
);
}
}
}
PhysicalWeightLayout::QuantizedBlockGrid {
packed_values,
packed_dimensions,
scales,
block_axes,
} => {
if !matches!(
logical_element_type,
ElementType::F16 | ElementType::Bf16 | ElementType::F32
) {
return Err(self.invalid(
"block-grid quantized logical weight dtype must be floating point",
));
}
let axes = [block_axes[0] as usize, block_axes[1] as usize];
if axes[0] >= axes[1] || axes[1] >= semantic_dimensions.len() {
return Err(self.invalid(
"block-grid quantization axes must be distinct, in range, and ascending",
));
}
let quantization = {
let component = self.component(&packed_values.component_id)?;
let WeightEncoding::Quantized(spec) = &component.encoding else {
return Err(self.invalid(
"block-grid packed-values component does not carry a quantization spec",
));
};
spec.clone()
};
let block_shape = quantization.grouping.block_shape_2d().ok_or_else(|| {
self.invalid(
"block-grid packed-values quantization spec must carry a two-dimensional block shape",
)
})?;
if quantization.zero_point_type.is_some() {
return Err(self.invalid(
"block-grid quantization does not support an implicit zero-point component",
));
}
let packed_bytes = checked_elements(semantic_dimensions)
.and_then(|elements| {
elements.checked_mul(u64::from(quantization.bits_per_weight))
})
.and_then(|bits| bits.checked_add(7))
.map(|bits| bits / 8)
.ok_or_else(|| self.invalid("block-grid packed-values size overflows u64"))?;
if packed_dimensions.is_empty()
|| checked_elements(packed_dimensions) != Some(packed_bytes)
{
return Err(self.invalid(format!(
"block-grid packed-values semantic shape contains {} storage bytes, expected {packed_bytes}",
checked_elements(packed_dimensions).map_or_else(
|| "an invalid or overflowing number of".to_owned(),
|value| value.to_string(),
)
)));
}
let packed = self.bind_component(
packed_values,
packed_dimensions,
WeightComponentRole::PackedValues,
depth,
)?;
if packed.encoding != WeightEncoding::Quantized(quantization.clone()) {
return Err(self.invalid(
"packed-values encoding changed while validating the block-grid tree",
));
}
let mut scale_dimensions = semantic_dimensions.to_vec();
for (axis, block_size) in axes.into_iter().zip(block_shape) {
scale_dimensions[axis] =
semantic_dimensions[axis].div_ceil(u64::from(block_size.get()));
}
let scales_component = self.bind_component(
scales,
&scale_dimensions,
WeightComponentRole::Scales,
depth,
)?;
if scales_component.dense_element_type() != Some(quantization.scale_type) {
return Err(self.invalid(
"block-grid scale component dtype differs from the quantization spec",
));
}
}
PhysicalWeightLayout::BlockQuantized {
blocks,
block_axis,
block_padding,
} => {
if !matches!(
logical_element_type,
ElementType::F16 | ElementType::Bf16 | ElementType::F32
) {
return Err(
self.invalid("block-quantized logical weight dtype must be floating point")
);
}
let axis = *block_axis as usize;
if axis >= semantic_dimensions.len() {
return Err(self.invalid("block quantization axis is out of range"));
}
let quantization = {
let component = self.component(&blocks.component_id)?;
let WeightEncoding::BlockQuantized(spec) = &component.encoding else {
return Err(self
.invalid("block component does not carry a block quantization spec"));
};
spec.clone()
};
let mut block_dimensions = self.grouped_dimensions(
semantic_dimensions,
block_padding,
axis,
u64::from(quantization.logical_values_per_block),
)?;
block_dimensions[axis] /= u64::from(quantization.logical_values_per_block);
let component = self.bind_component(
blocks,
&block_dimensions,
WeightComponentRole::PackedValues,
depth,
)?;
if component.encoding != WeightEncoding::BlockQuantized(quantization) {
return Err(
self.invalid("block encoding changed while validating the physical layout")
);
}
}
PhysicalWeightLayout::AxisReshapePermutation {
values,
axis,
logical_offset,
extent,
reshape,
stored_axis_order,
} => {
let axis = *axis as usize;
let end = logical_offset.checked_add(*extent);
let reshape_rank = reshape.len();
let order_is_permutation = stored_axis_order.len() == reshape_rank
&& stored_axis_order
.iter()
.all(|axis| (*axis as usize) < reshape_rank)
&& stored_axis_order
.iter()
.copied()
.collect::<BTreeSet<_>>()
.len()
== reshape_rank;
let order_is_identity = stored_axis_order
.iter()
.filter(|stored| reshape[**stored as usize] > 1)
.copied()
.eq(reshape
.iter()
.enumerate()
.filter_map(|(axis, extent)| (*extent > 1).then_some(axis as u32)));
if axis >= semantic_dimensions.len()
|| *extent == 0
|| end.is_none_or(|end| end > semantic_dimensions[axis])
|| reshape_rank < 2
|| reshape.iter().any(|dimension| *dimension == 0)
|| checked_elements(reshape) != Some(*extent)
|| !order_is_permutation
|| order_is_identity
{
return Err(self.invalid(
"axis reshape permutation has invalid range, shape, or stored axis order",
));
}
self.validate_layout(values, semantic_dimensions, logical_element_type, depth + 1)?;
}
PhysicalWeightLayout::Indexed {
indices,
values,
source_axis_extent,
} => {
let axis = indices.axis as usize;
if axis >= semantic_dimensions.len() || *source_axis_extent == 0 {
return Err(self.invalid("indexed layout axis or source extent is invalid"));
}
self.validate_axis_component(
indices,
semantic_dimensions,
axis,
WeightComponentRole::Indices,
true,
depth,
)?;
let mut source_dimensions = semantic_dimensions.to_vec();
source_dimensions[axis] = *source_axis_extent;
checked_elements(&source_dimensions)
.ok_or_else(|| self.invalid("indexed source semantic shape overflows u64"))?;
self.validate_layout(values, &source_dimensions, logical_element_type, depth + 1)?;
}
PhysicalWeightLayout::ExpertStack {
experts,
expert_axis,
} => {
let axis = *expert_axis as usize;
if axis >= semantic_dimensions.len() {
return Err(self.invalid("expert stack axis is out of range"));
}
let expected_count = usize::try_from(semantic_dimensions[axis]).map_err(|_| {
self.invalid("expert stack count does not fit the platform usize")
})?;
if experts.is_empty() || experts.len() != expected_count {
return Err(self
.invalid("expert stack child count differs from its logical expert axis"));
}
let mut expert_dimensions = semantic_dimensions.to_vec();
expert_dimensions.remove(axis);
if expert_dimensions.is_empty() {
return Err(self.invalid(
"expert stack children must retain at least one tensor dimension",
));
}
for expert in experts {
self.validate_layout(
expert,
&expert_dimensions,
logical_element_type,
depth + 1,
)?;
}
}
}
Ok(())
}
}
fn is_axis_permutation(axis_order: &[u32], rank: usize) -> bool {
axis_order.len() == rank
&& axis_order.iter().all(|axis| (*axis as usize) < rank)
&& axis_order.iter().copied().collect::<BTreeSet<_>>().len() == rank
}
fn checked_round_up(extent: u64, multiple: u64) -> Option<u64> {
if extent == 0 || multiple == 0 {
return None;
}
extent
.checked_add(multiple.checked_sub(1)?)
.map(|rounded| rounded / multiple * multiple)
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProgramTensorSpec {
pub dimensions: Vec<u64>,
pub element_type: ElementType,
pub layout: ResolvedTensorLayout,
}
impl ProgramTensorSpec {
pub fn validate(&self, field: &str) -> Result<(), VNextError> {
super::ResolvedTensorSpec::new(
self.dimensions.clone(),
self.element_type,
self.layout.clone(),
)
.map(|_| ())
.map_err(|error| VNextError::InvalidExecutionPlan {
reason: format!("{field} is invalid: {error}"),
})
}
pub fn byte_len(&self) -> Result<u64, VNextError> {
checked_elements(&self.dimensions)
.and_then(|elements| elements.checked_mul(self.element_type.size_bytes()))
.ok_or_else(|| VNextError::InvalidExecutionPlan {
reason: "program tensor byte size overflows u64".to_owned(),
})
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct WeightReference {
pub weight_id: WeightId,
pub value_id: ProgramValueId,
pub tensor: ProgramTensorSpec,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct StateSpec {
pub id: StateId,
pub value_id: ProgramValueId,
pub tensor: ProgramTensorSpec,
pub lifetime: StateLifetime,
pub capacity_demand: StateCapacityDemand,
pub initialization: StateInitialization,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct StateSpecWire {
id: StateId,
value_id: ProgramValueId,
tensor: ProgramTensorSpec,
lifetime: StateLifetime,
capacity_demand: StateCapacityDemand,
initialization: StateInitialization,
}
impl<'de> Deserialize<'de> for StateSpec {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let wire = StateSpecWire::deserialize(deserializer)?;
wire.tensor
.validate("state_spec.tensor")
.and_then(|()| wire.capacity_demand.validate(wire.tensor.byte_len()?))
.map_err(serde::de::Error::custom)?;
Ok(Self {
id: wire.id,
value_id: wire.value_id,
tensor: wire.tensor,
lifetime: wire.lifetime,
capacity_demand: wire.capacity_demand,
initialization: wire.initialization,
})
}
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum StateLifetime {
Request,
Sequence,
Step,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum StateCapacityDemand {
FixedPerScope,
TokenScaled {
bytes_per_token: u64,
maximum_tokens: u64,
},
}
#[derive(Deserialize)]
#[serde(rename_all = "snake_case", deny_unknown_fields)]
enum StateCapacityDemandWire {
FixedPerScope,
TokenScaled {
bytes_per_token: u64,
maximum_tokens: u64,
},
}
impl<'de> Deserialize<'de> for StateCapacityDemand {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let demand = match StateCapacityDemandWire::deserialize(deserializer)? {
StateCapacityDemandWire::FixedPerScope => Self::FixedPerScope,
StateCapacityDemandWire::TokenScaled {
bytes_per_token,
maximum_tokens,
} => Self::TokenScaled {
bytes_per_token,
maximum_tokens,
},
};
demand.validate(1).map_err(serde::de::Error::custom)?;
Ok(demand)
}
}
impl StateCapacityDemand {
pub fn validate(self, tensor_minimum_bytes: u64) -> Result<(), VNextError> {
let valid = match self {
Self::FixedPerScope => tensor_minimum_bytes > 0,
Self::TokenScaled {
bytes_per_token,
maximum_tokens,
} => {
bytes_per_token >= tensor_minimum_bytes
&& maximum_tokens > 0
&& bytes_per_token.checked_mul(maximum_tokens).is_some()
}
};
if !valid {
return Err(VNextError::InvalidExecutionPlan {
reason: "state resource demand is zero, smaller than its tensor, or overflows u64"
.to_owned(),
});
}
Ok(())
}
pub fn minimum_bytes(self, tensor_minimum_bytes: u64) -> Result<u64, VNextError> {
self.validate(tensor_minimum_bytes)?;
Ok(match self {
Self::FixedPerScope => tensor_minimum_bytes,
Self::TokenScaled {
bytes_per_token, ..
} => bytes_per_token,
})
}
pub fn theoretical_bytes(self, tensor_minimum_bytes: u64) -> Result<u64, VNextError> {
self.validate(tensor_minimum_bytes)?;
match self {
Self::FixedPerScope => Ok(tensor_minimum_bytes),
Self::TokenScaled {
bytes_per_token,
maximum_tokens,
} => bytes_per_token.checked_mul(maximum_tokens).ok_or_else(|| {
VNextError::InvalidExecutionPlan {
reason: "token-scaled state demand overflows u64".to_owned(),
}
}),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum ProgramNodeWorkSpec {
Fixed,
Tokens { value_id: ProgramValueId, axis: u32 },
}
impl ProgramNodeWorkSpec {
pub fn tokens(value_id: ProgramValueId, axis: u32) -> Self {
Self::Tokens { value_id, axis }
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProgramNode {
pub id: NodeId,
pub operation_id: OperationId,
pub required_version: ContractVersion,
pub work: ProgramNodeWorkSpec,
pub inputs: Vec<ProgramValueId>,
pub outputs: Vec<ProgramValueId>,
pub attributes: BTreeMap<AttributeId, SemanticValue>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ProgramBlock {
pub id: String,
pub nodes: Vec<ProgramNode>,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct ModelProgram {
family_id: ModelFamilyId,
inputs: Vec<ProgramValueId>,
blocks: Vec<ProgramBlock>,
states: Vec<StateSpec>,
weights: Vec<WeightReference>,
outputs: Vec<ProgramValueId>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct ModelProgramWire {
family_id: ModelFamilyId,
inputs: Vec<ProgramValueId>,
blocks: Vec<ProgramBlock>,
states: Vec<StateSpec>,
weights: Vec<WeightReference>,
outputs: Vec<ProgramValueId>,
}
impl<'de> Deserialize<'de> for ModelProgram {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let wire = ModelProgramWire::deserialize(deserializer)?;
Self::new(
wire.family_id,
wire.inputs,
wire.blocks,
wire.states,
wire.weights,
wire.outputs,
)
.map_err(serde::de::Error::custom)
}
}
impl ModelProgram {
pub fn new(
family_id: ModelFamilyId,
inputs: Vec<ProgramValueId>,
blocks: Vec<ProgramBlock>,
mut states: Vec<StateSpec>,
mut weights: Vec<WeightReference>,
outputs: Vec<ProgramValueId>,
) -> Result<Self, VNextError> {
if blocks.is_empty() {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.blocks".to_owned(),
reason: "at least one block is required".to_owned(),
});
}
let mut known_values = BTreeSet::new();
if inputs.is_empty()
|| inputs
.iter()
.any(|input| !known_values.insert(input.clone()))
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.inputs".to_owned(),
reason: "input identities must be non-empty and unique".to_owned(),
});
}
let mut block_ids = BTreeSet::new();
let mut node_ids = BTreeSet::new();
for state in &states {
let tensor_valid = state
.tensor
.validate(&format!("program.states.{}.tensor", state.id))
.and_then(|()| state.capacity_demand.validate(state.tensor.byte_len()?));
if tensor_valid.is_err() || !known_values.insert(state.value_id.clone()) {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.states.value_id".to_owned(),
reason: format!("duplicate value `{}`", state.value_id),
});
}
}
let mut weight_ids = BTreeSet::new();
for weight in &weights {
if !weight_ids.insert(weight.weight_id.clone())
|| !known_values.insert(weight.value_id.clone())
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.weights".to_owned(),
reason: format!(
"duplicate weight `{}` or value `{}`",
weight.weight_id, weight.value_id
),
});
}
weight
.tensor
.validate(&format!("program.weights.{}.tensor", weight.weight_id))?;
}
for block in &blocks {
if block.id.is_empty() || block.nodes.is_empty() || !block_ids.insert(block.id.clone())
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.blocks.id".to_owned(),
reason: "block identities must be non-empty and unique".to_owned(),
});
}
for node in &block.nodes {
if node.required_version.major == 0 || node.outputs.is_empty() {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.nodes.contract".to_owned(),
reason: format!("node `{}` has an invalid version or no outputs", node.id),
});
}
for value in node.attributes.values() {
value.validate(&format!("program node `{}` attributes", node.id))?;
}
if !node_ids.insert(node.id.clone()) {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.nodes.id".to_owned(),
reason: format!("duplicate node `{}`", node.id),
});
}
if let ProgramNodeWorkSpec::Tokens { value_id, .. } = &node.work {
let source_count = node
.inputs
.iter()
.chain(&node.outputs)
.filter(|candidate| *candidate == value_id)
.count();
let is_state_or_weight = states.iter().any(|state| state.value_id == *value_id)
|| weights.iter().any(|weight| weight.value_id == *value_id);
if source_count != 1 || is_state_or_weight {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.nodes.work".to_owned(),
reason: format!(
"node `{}` token work source must identify one activation binding",
node.id
),
});
}
}
if node
.inputs
.iter()
.any(|input| !known_values.contains(input))
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.nodes.inputs".to_owned(),
reason: format!("node `{}` references an unknown input", node.id),
});
}
for output in &node.outputs {
if !known_values.insert(output.clone()) {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.nodes.outputs".to_owned(),
reason: format!("value `{output}` has multiple producers"),
});
}
}
}
}
let mut state_ids = BTreeSet::new();
if states
.iter()
.any(|state| !state_ids.insert(state.id.clone()))
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.states.id".to_owned(),
reason: "state identities must be unique".to_owned(),
});
}
let mut output_ids = BTreeSet::new();
if outputs.is_empty()
|| outputs
.iter()
.any(|output| !known_values.contains(output) || !output_ids.insert(output.clone()))
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.outputs".to_owned(),
reason: "program outputs must be non-empty, known, and unique".to_owned(),
});
}
states.sort_by(|left, right| left.id.cmp(&right.id));
weights.sort_by(|left, right| left.weight_id.cmp(&right.weight_id));
Ok(Self {
family_id,
inputs,
blocks,
states,
weights,
outputs,
})
}
pub fn family_id(&self) -> &ModelFamilyId {
&self.family_id
}
pub fn inputs(&self) -> &[ProgramValueId] {
&self.inputs
}
pub fn blocks(&self) -> &[ProgramBlock] {
&self.blocks
}
pub fn states(&self) -> &[StateSpec] {
&self.states
}
pub fn weights(&self) -> &[WeightReference] {
&self.weights
}
pub fn outputs(&self) -> &[ProgramValueId] {
&self.outputs
}
pub fn fingerprint(&self) -> Result<String, VNextError> {
let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
context: "serialize model program",
message: error.to_string(),
})?;
Ok(format!("{:x}", Sha256::digest(bytes)))
}
}
impl WeightSchema {
pub fn validate_program_references(
&self,
family_id: &ModelFamilyId,
program: &ModelProgram,
) -> Result<(), VNextError> {
if program.family_id() != family_id {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.family_id".to_owned(),
reason: "program family does not match the weight schema owner".to_owned(),
});
}
let schema_weights = self
.tensors
.iter()
.map(|tensor| (&tensor.id, tensor.required))
.collect::<BTreeMap<_, _>>();
let referenced_weights = program
.weights()
.iter()
.map(|reference| &reference.weight_id)
.collect::<BTreeSet<_>>();
if let Some(weight_id) = referenced_weights
.iter()
.find(|weight_id| !schema_weights.contains_key(**weight_id))
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.weights".to_owned(),
reason: format!("program references unknown weight `{weight_id}`"),
});
}
if let Some(weight_id) = schema_weights.iter().find_map(|(weight_id, required)| {
(*required && !referenced_weights.contains(weight_id)).then_some(*weight_id)
}) {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.weights".to_owned(),
reason: format!("program does not reference required weight `{weight_id}`"),
});
}
for reference in program.weights() {
let tensor = self.tensor(&reference.weight_id).ok_or_else(|| {
VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.weights".to_owned(),
reason: format!(
"program references unknown weight `{}`",
reference.weight_id
),
}
})?;
if reference.tensor.dimensions != tensor.dimensions
|| reference.tensor.element_type != tensor.logical_element_type
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: format!("program.weights.{}.tensor", reference.weight_id),
reason: "program value shape or dtype differs from the logical weight schema"
.to_owned(),
});
}
}
Ok(())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TemplateMetadata {
pub template: String,
pub source_file: String,
pub sha256: String,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SpecialTokenRole {
Bos,
Eos,
Pad,
Stop,
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Serialize)]
pub struct SpecialTokenCollision {
first: SpecialTokenRole,
second: SpecialTokenRole,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct SpecialTokenCollisionWire {
first: SpecialTokenRole,
second: SpecialTokenRole,
}
impl<'de> Deserialize<'de> for SpecialTokenCollision {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let wire = SpecialTokenCollisionWire::deserialize(deserializer)?;
Self::new(wire.first, wire.second).map_err(serde::de::Error::custom)
}
}
impl SpecialTokenCollision {
pub fn new(first: SpecialTokenRole, second: SpecialTokenRole) -> Result<Self, VNextError> {
if first == second {
return Err(VNextError::InvalidExecutionPlan {
reason: "a special-token collision must name two different roles".to_owned(),
});
}
let (first, second) = if first < second {
(first, second)
} else {
(second, first)
};
Ok(Self { first, second })
}
pub const fn first(&self) -> SpecialTokenRole {
self.first
}
pub const fn second(&self) -> SpecialTokenRole {
self.second
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct SpecialTokenCollisionPolicy {
allowed: BTreeSet<SpecialTokenCollision>,
}
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct SpecialTokenCollisionPolicyWire {
allowed: BTreeSet<SpecialTokenCollision>,
}
impl<'de> Deserialize<'de> for SpecialTokenCollisionPolicy {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let wire = SpecialTokenCollisionPolicyWire::deserialize(deserializer)?;
Ok(Self::new(wire.allowed))
}
}
impl SpecialTokenCollisionPolicy {
pub fn new(allowed: BTreeSet<SpecialTokenCollision>) -> Self {
Self { allowed }
}
pub fn require_distinct() -> Self {
Self {
allowed: BTreeSet::new(),
}
}
pub fn allows(&self, left: SpecialTokenRole, right: SpecialTokenRole) -> bool {
SpecialTokenCollision::new(left, right)
.is_ok_and(|collision| self.allowed.contains(&collision))
}
pub fn allowed(&self) -> &BTreeSet<SpecialTokenCollision> {
&self.allowed
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct SpecialTokenMetadata {
pub bos_token_id: Option<u32>,
pub eos_token_ids: BTreeSet<u32>,
pub pad_token_id: Option<u32>,
pub collision_policy: SpecialTokenCollisionPolicy,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ModelSemanticMetadata {
pub template: TemplateMetadata,
pub special_tokens: SpecialTokenMetadata,
}
pub trait ModelFamilyProvider: Send + Sync {
type Config: Clone + Send + Sync + Serialize + DeserializeOwned + 'static;
fn family_id(&self) -> &ModelFamilyId;
fn external_metadata_ids(&self) -> BTreeSet<ExternalModelMetadataId>;
fn validate_config_identity(
&self,
raw: &serde_json::Value,
config: &Self::Config,
) -> Result<(), VNextError>;
fn validated_external_metadata_id(
&self,
raw: &serde_json::Value,
config: &Self::Config,
) -> Result<ExternalModelMetadataId, VNextError>;
fn parse_config(&self, raw: &serde_json::Value) -> Result<Self::Config, VNextError>;
fn weight_schema(&self, config: &Self::Config) -> Result<WeightSchema, VNextError>;
fn semantic_program(&self, config: &Self::Config) -> Result<ModelProgram, VNextError>;
fn semantic_metadata(&self, config: &Self::Config)
-> Result<ModelSemanticMetadata, VNextError>;
}
pub const MAX_PREPARED_MODEL_FAMILY_WIRE_BYTES: usize = 16 * 1024 * 1024;
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct PreparedModelFamily {
family_id: ModelFamilyId,
external_metadata_id: ExternalModelMetadataId,
canonical_config: serde_json::Value,
config_fingerprint: String,
weight_schema: WeightSchema,
program: ModelProgram,
metadata: ModelSemanticMetadata,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub struct UnvalidatedPreparedModelFamily {
family_id: ModelFamilyId,
external_metadata_id: ExternalModelMetadataId,
canonical_config: serde_json::Value,
config_fingerprint: String,
weight_schema: WeightSchema,
program: ModelProgram,
metadata: ModelSemanticMetadata,
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize)]
pub(crate) struct PreparedModelFamilyWire {
family_id: ModelFamilyId,
external_metadata_id: ExternalModelMetadataId,
canonical_config: serde_json::Value,
config_fingerprint: String,
weight_schema: WeightSchema,
program: ModelProgram,
metadata: ModelSemanticMetadata,
}
#[derive(Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
struct PreparedModelFamilyWireFields {
family_id: ModelFamilyId,
external_metadata_id: ExternalModelMetadataId,
canonical_config: serde_json::Value,
config_fingerprint: String,
weight_schema: WeightSchema,
program: ModelProgram,
metadata: ModelSemanticMetadata,
}
impl<'de> Deserialize<'de> for PreparedModelFamilyWire {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let raw = serde_json::Value::deserialize(deserializer)?;
let fields =
PreparedModelFamilyWireFields::deserialize(&raw).map_err(serde::de::Error::custom)?;
let canonical = serde_json::to_value(&fields).map_err(serde::de::Error::custom)?;
if canonical != raw {
return Err(serde::de::Error::custom(
"prepared model family wire contains unknown or non-canonical nested fields",
));
}
Ok(Self {
family_id: fields.family_id,
external_metadata_id: fields.external_metadata_id,
canonical_config: fields.canonical_config,
config_fingerprint: fields.config_fingerprint,
weight_schema: fields.weight_schema,
program: fields.program,
metadata: fields.metadata,
})
}
}
impl From<PreparedModelFamilyWire> for UnvalidatedPreparedModelFamily {
fn from(wire: PreparedModelFamilyWire) -> Self {
Self {
family_id: wire.family_id,
external_metadata_id: wire.external_metadata_id,
canonical_config: wire.canonical_config,
config_fingerprint: wire.config_fingerprint,
weight_schema: wire.weight_schema,
program: wire.program,
metadata: wire.metadata,
}
}
}
impl UnvalidatedPreparedModelFamily {
pub fn revalidate(
self,
registry: &dyn ModelFamilyRegistry,
) -> Result<PreparedModelFamily, VNextError> {
let registration = registry.resolve(&self.family_id)?;
if registration.family_id() != &self.family_id {
return Err(VNextError::InvalidModelConfig {
family_id: self.family_id.to_string(),
field: "registration.family_id".to_owned(),
reason: "registry returned a registration for a different family".to_owned(),
});
}
let metadata_registration = registry.resolve_external(&self.external_metadata_id)?;
if !std::ptr::eq(registration, metadata_registration) {
return Err(VNextError::InvalidModelConfig {
family_id: self.family_id.to_string(),
field: "external_metadata_id".to_owned(),
reason: "external metadata identity resolves to a different family registration"
.to_owned(),
});
}
let rebuilt = registration.prepare(&self.canonical_config)?;
let exact_match = rebuilt.family_id == self.family_id
&& rebuilt.external_metadata_id == self.external_metadata_id
&& rebuilt.canonical_config == self.canonical_config
&& rebuilt.config_fingerprint == self.config_fingerprint
&& rebuilt.weight_schema == self.weight_schema
&& rebuilt.program == self.program
&& rebuilt.metadata == self.metadata;
if !exact_match {
return Err(VNextError::InvalidModelConfig {
family_id: self.family_id.to_string(),
field: "prepared_package".to_owned(),
reason: "serialized package differs from the typed provider reconstruction"
.to_owned(),
});
}
Ok(rebuilt)
}
}
impl PreparedModelFamily {
fn from_canonical_config(
family_id: ModelFamilyId,
external_metadata_id: ExternalModelMetadataId,
canonical_config: serde_json::Value,
mut weight_schema: WeightSchema,
program: ModelProgram,
metadata: ModelSemanticMetadata,
) -> Result<Self, VNextError> {
if !canonical_config.is_object()
|| canonicalize_json(canonical_config.clone()) != canonical_config
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "config".to_owned(),
reason: "prepared config must be a canonical JSON object".to_owned(),
});
}
let config_bytes =
serde_json::to_vec(&canonical_config).map_err(|error| VNextError::Serialization {
context: "serialize canonical model family config",
message: error.to_string(),
})?;
let config_fingerprint = format!("{:x}", Sha256::digest(config_bytes));
weight_schema.validate(&family_id)?;
weight_schema.normalize();
weight_schema.validate(&family_id)?;
if program.family_id() != &family_id {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "program.family_id".to_owned(),
reason: "program family does not match prepared family".to_owned(),
});
}
weight_schema.validate_program_references(&family_id, &program)?;
Self::validate_metadata(&family_id, &metadata)?;
Ok(Self {
family_id,
external_metadata_id,
canonical_config,
config_fingerprint,
weight_schema,
program,
metadata,
})
}
fn validate_metadata(
family_id: &ModelFamilyId,
metadata: &ModelSemanticMetadata,
) -> Result<(), VNextError> {
let source = metadata.template.source_file.as_str();
let valid_source = !source.is_empty()
&& !source.starts_with('/')
&& !source.contains('\\')
&& source
.split('/')
.all(|component| !matches!(component, "" | "." | ".."));
if metadata.template.template.is_empty()
|| !valid_source
|| !is_canonical_sha256(&metadata.template.sha256)
|| metadata.special_tokens.eos_token_ids.is_empty()
{
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "semantic_metadata".to_owned(),
reason: "template, source, checksum, and end tokens must be explicit and valid"
.to_owned(),
});
}
Ok(())
}
pub fn family_id(&self) -> &ModelFamilyId {
&self.family_id
}
pub fn external_metadata_id(&self) -> &ExternalModelMetadataId {
&self.external_metadata_id
}
pub fn canonical_config(&self) -> &serde_json::Value {
&self.canonical_config
}
pub fn config_fingerprint(&self) -> &str {
&self.config_fingerprint
}
pub fn weight_schema(&self) -> &WeightSchema {
&self.weight_schema
}
pub fn program(&self) -> &ModelProgram {
&self.program
}
pub fn metadata(&self) -> &ModelSemanticMetadata {
&self.metadata
}
pub fn fingerprint(&self) -> Result<String, VNextError> {
let bytes = serde_json::to_vec(self).map_err(|error| VNextError::Serialization {
context: "serialize prepared model family",
message: error.to_string(),
})?;
Ok(format!("{:x}", Sha256::digest(bytes)))
}
pub fn decode_untrusted(bytes: &[u8]) -> Result<UnvalidatedPreparedModelFamily, VNextError> {
if bytes.len() > MAX_PREPARED_MODEL_FAMILY_WIRE_BYTES {
return Err(VNextError::Serialization {
context: "decode untrusted prepared model family",
message: format!(
"payload has {} bytes; maximum is {MAX_PREPARED_MODEL_FAMILY_WIRE_BYTES}",
bytes.len()
),
});
}
serde_json::from_slice::<PreparedModelFamilyWire>(bytes)
.map(Into::into)
.map_err(|error| VNextError::Serialization {
context: "decode untrusted prepared model family",
message: error.to_string(),
})
}
pub fn from_json_validated(
bytes: &[u8],
registry: &dyn ModelFamilyRegistry,
) -> Result<Self, VNextError> {
Self::decode_untrusted(bytes)?.revalidate(registry)
}
}
fn canonicalize_json(value: serde_json::Value) -> serde_json::Value {
match value {
serde_json::Value::Array(values) => {
serde_json::Value::Array(values.into_iter().map(canonicalize_json).collect())
}
serde_json::Value::Object(values) => {
let sorted = values
.into_iter()
.map(|(key, value)| (key, canonicalize_json(value)))
.collect::<BTreeMap<_, _>>();
serde_json::Value::Object(sorted.into_iter().collect())
}
other => other,
}
}
fn validate_raw_config_consumed(
family_id: &ModelFamilyId,
raw: &serde_json::Value,
typed: &serde_json::Value,
) -> Result<(), VNextError> {
fn walk(raw: &serde_json::Value, typed: &serde_json::Value, path: &str) -> Option<String> {
match (raw, typed) {
(serde_json::Value::Object(raw), serde_json::Value::Object(typed)) => {
for (key, raw_value) in raw {
let next = if path.is_empty() {
format!("/{key}")
} else {
format!("{path}/{key}")
};
let Some(typed_value) = typed.get(key) else {
return Some(next);
};
if let Some(rejected) = walk(raw_value, typed_value, &next) {
return Some(rejected);
}
}
None
}
(serde_json::Value::Array(raw), serde_json::Value::Array(typed))
if raw.len() == typed.len() =>
{
raw.iter()
.zip(typed)
.enumerate()
.find_map(|(index, (raw, typed))| walk(raw, typed, &format!("{path}/{index}")))
}
_ if raw == typed => None,
_ => Some(path.to_owned()),
}
}
if !raw.is_object() || !typed.is_object() {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: "config".to_owned(),
reason: "raw and typed model configurations must be JSON objects".to_owned(),
});
}
if let Some(path) = walk(raw, typed, "") {
return Err(VNextError::InvalidModelConfig {
family_id: family_id.to_string(),
field: if path.is_empty() {
"config".to_owned()
} else {
path
},
reason: "raw configuration field was ignored or changed by typed parsing".to_owned(),
});
}
Ok(())
}
fn is_canonical_sha256(value: &str) -> bool {
value.len() == 64
&& value
.bytes()
.all(|byte| byte.is_ascii_digit() || (b'a'..=b'f').contains(&byte))
}
pub trait ModelFamilyRegistration: Send + Sync {
fn family_id(&self) -> &ModelFamilyId;
fn external_metadata_ids(&self) -> BTreeSet<ExternalModelMetadataId>;
fn prepare(&self, raw_config: &serde_json::Value) -> Result<PreparedModelFamily, VNextError>;
}
pub struct TypedFamilyRegistration<P> {
provider: P,
}
impl<P> TypedFamilyRegistration<P> {
pub fn new(provider: P) -> Self {
Self { provider }
}
}
impl<P: ModelFamilyProvider> ModelFamilyRegistration for TypedFamilyRegistration<P> {
fn family_id(&self) -> &ModelFamilyId {
self.provider.family_id()
}
fn external_metadata_ids(&self) -> BTreeSet<ExternalModelMetadataId> {
self.provider.external_metadata_ids()
}
fn prepare(&self, raw_config: &serde_json::Value) -> Result<PreparedModelFamily, VNextError> {
let external_metadata_ids = self.provider.external_metadata_ids();
if external_metadata_ids.is_empty() {
return Err(VNextError::InvalidModelConfig {
family_id: self.provider.family_id().to_string(),
field: "external_metadata_ids".to_owned(),
reason: "model family must declare at least one external metadata identity"
.to_owned(),
});
}
let config = self.provider.parse_config(raw_config)?;
let external_metadata_id = self
.provider
.validated_external_metadata_id(raw_config, &config)?;
if !external_metadata_ids.contains(&external_metadata_id) {
return Err(VNextError::InvalidModelConfig {
family_id: self.provider.family_id().to_string(),
field: "external_metadata_id".to_owned(),
reason: format!(
"provider selected undeclared external metadata identity `{external_metadata_id}`"
),
});
}
let typed_config = canonicalize_json(serde_json::to_value(&config).map_err(|error| {
VNextError::Serialization {
context: "serialize typed model configuration",
message: error.to_string(),
}
})?);
validate_raw_config_consumed(self.provider.family_id(), raw_config, &typed_config)?;
let weight_schema = self.provider.weight_schema(&config)?;
let program = self.provider.semantic_program(&config)?;
let metadata = self.provider.semantic_metadata(&config)?;
PreparedModelFamily::from_canonical_config(
self.provider.family_id().clone(),
external_metadata_id,
typed_config,
weight_schema,
program,
metadata,
)
}
}
pub trait ModelFamilyRegistry: Send + Sync {
fn registrations(&self) -> Vec<&dyn ModelFamilyRegistration>;
}
impl dyn ModelFamilyRegistry + '_ {
pub fn resolve(
&self,
family_id: &ModelFamilyId,
) -> Result<&dyn ModelFamilyRegistration, VNextError> {
let matches = self
.registrations()
.into_iter()
.filter(|registration| registration.family_id() == family_id)
.collect::<Vec<_>>();
match matches.as_slice() {
[] => Err(VNextError::UnknownModelFamily {
family_id: family_id.to_string(),
}),
[registration] => Ok(*registration),
_ => Err(VNextError::AmbiguousModelFamilyRegistration {
identity_kind: "internal family",
identity: family_id.to_string(),
matches: matches.len(),
}),
}
}
pub fn resolve_external(
&self,
metadata_id: &ExternalModelMetadataId,
) -> Result<&dyn ModelFamilyRegistration, VNextError> {
let matches = self
.registrations()
.into_iter()
.filter(|registration| registration.external_metadata_ids().contains(metadata_id))
.collect::<Vec<_>>();
match matches.as_slice() {
[] => Err(VNextError::UnknownExternalModelMetadata {
metadata_id: metadata_id.to_string(),
}),
[registration] => Ok(*registration),
_ => Err(VNextError::AmbiguousModelFamilyRegistration {
identity_kind: "external metadata",
identity: metadata_id.to_string(),
matches: matches.len(),
}),
}
}
}
#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct TokenizerDescriptor {
pub tokenizer_id: TokenizerId,
pub source_file: String,
pub sha256: String,
pub vocabulary_size: u64,
}