use std::collections::BTreeSet;
use std::fmt;
use serde_json::Value;
use crate::census::{CensusReport, WeightsManifest};
use crate::fttsq::{
AccessClass, FttsqError, FttsqStreamPlan, FttsqStreamingWriter, StoredDtype,
TensorEntry as ArtifactTensorEntry,
};
use crate::safetensors::{Dtype, SafetensorsIndex, TensorView, WeightsError};
use crate::sha256::Sha256;
pub const MAX_Q8_OUTPUT_CHANNEL_WIDTH: usize = 65_536;
pub const MAX_Q8_OUTPUT_CHANNELS: usize = 262_144;
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum TensorStoragePolicy {
Verbatim,
Q8PerOutputChannel,
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub struct TensorConversion {
source_name: String,
artifact_name: String,
access_class: AccessClass,
storage: TensorStoragePolicy,
}
impl TensorConversion {
#[must_use]
pub fn verbatim(
source_name: impl Into<String>,
artifact_name: impl Into<String>,
access_class: AccessClass,
) -> Self {
Self {
source_name: source_name.into(),
artifact_name: artifact_name.into(),
access_class,
storage: TensorStoragePolicy::Verbatim,
}
}
#[must_use]
pub fn q8_per_output_channel(
source_name: impl Into<String>,
artifact_name: impl Into<String>,
access_class: AccessClass,
) -> Self {
Self {
source_name: source_name.into(),
artifact_name: artifact_name.into(),
access_class,
storage: TensorStoragePolicy::Q8PerOutputChannel,
}
}
fn section_name(&self) -> &'static str {
self.access_class.as_str()
}
fn scales_name(&self) -> String {
format!("{}.scales", self.artifact_name)
}
}
#[derive(Clone, Debug)]
pub struct StreamingConversionPlan {
model_family: String,
source_sha256: String,
license_notice: String,
model_config: Value,
quantization_manifest: Value,
tensors: Vec<TensorConversion>,
}
impl StreamingConversionPlan {
#[must_use]
pub fn new(model_family: impl Into<String>, source_sha256: impl Into<String>) -> Self {
Self {
model_family: model_family.into(),
source_sha256: source_sha256.into(),
license_notice: String::new(),
model_config: Value::Null,
quantization_manifest: Value::Null,
tensors: Vec::new(),
}
}
#[must_use]
pub fn license_notice(mut self, notice: impl Into<String>) -> Self {
self.license_notice = notice.into();
self
}
#[must_use]
pub fn model_config(mut self, config: Value) -> Self {
self.model_config = config;
self
}
#[must_use]
pub fn quantization_manifest(mut self, manifest: Value) -> Self {
self.quantization_manifest = manifest;
self
}
#[must_use]
pub fn tensor(mut self, tensor: TensorConversion) -> Self {
self.tensors.push(tensor);
self
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum ConversionPlanError {
NoTensorPolicies,
DuplicateSourcePolicy {
name: String,
},
SourceTensorMissing {
name: String,
},
SourceTensorUnplanned {
name: String,
},
DuplicateArtifactTensor {
name: String,
},
EmptyArtifactTensorName {
source_name: String,
},
Q8RequiresMatrix {
name: String,
rank: usize,
},
Q8EmptyOutputChannel {
name: String,
},
Q8OutputChannelTooWide {
name: String,
width: usize,
limit: usize,
},
Q8OutputChannelCountTooLarge {
name: String,
rows: usize,
limit: usize,
},
ShapeOutOfRange {
name: String,
},
SectionLengthOverflow {
name: String,
},
}
impl fmt::Display for ConversionPlanError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::NoTensorPolicies => f.write_str("conversion plan has no tensor policies"),
Self::DuplicateSourcePolicy { name } => {
write!(
f,
"conversion plan names source tensor `{name}` more than once"
)
}
Self::SourceTensorMissing { name } => {
write!(
f,
"conversion plan names source tensor `{name}`, which is absent"
)
}
Self::SourceTensorUnplanned { name } => {
write!(
f,
"source tensor `{name}` has no explicit conversion policy"
)
}
Self::DuplicateArtifactTensor { name } => {
write!(
f,
"conversion plan would emit artifact tensor `{name}` more than once"
)
}
Self::EmptyArtifactTensorName { source_name } => write!(
f,
"conversion plan gives source tensor `{source_name}` an empty artifact name"
),
Self::Q8RequiresMatrix { name, rank } => write!(
f,
"Q8 conversion for `{name}` requires rank 2 or greater, got rank {rank}"
),
Self::Q8EmptyOutputChannel { name } => {
write!(f, "Q8 conversion for `{name}` has an empty output channel")
}
Self::Q8OutputChannelTooWide { name, width, limit } => write!(
f,
"Q8 conversion for `{name}` has output-channel width {width}, exceeding {limit}"
),
Self::Q8OutputChannelCountTooLarge { name, rows, limit } => write!(
f,
"Q8 conversion for `{name}` has {rows} output channels, exceeding {limit}"
),
Self::ShapeOutOfRange { name } => {
write!(
f,
"source tensor `{name}` has a shape outside the artifact range"
)
}
Self::SectionLengthOverflow { name } => {
write!(
f,
"source tensor `{name}` overflows its planned artifact section length"
)
}
}
}
}
impl std::error::Error for ConversionPlanError {}
#[derive(Debug)]
pub enum StreamingConversionError {
Source(WeightsError),
SourceCensus(Box<CensusReport>),
SourceDigestMismatch {
expected: String,
actual: String,
},
Plan(ConversionPlanError),
Artifact(FttsqError),
Quantization(MatrixQuantizationError<Q8SectionSinkError>),
}
impl fmt::Display for StreamingConversionError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::Source(error) => write!(f, "cannot parse source checkpoint: {error}"),
Self::SourceCensus(report) => f.write_str(&report.render()),
Self::SourceDigestMismatch { expected, actual } => write!(
f,
"source checkpoint SHA-256 mismatch: expected {expected}, got {actual}"
),
Self::Plan(error) => write!(f, "invalid conversion plan: {error}"),
Self::Artifact(error) => write!(f, "cannot write .fttsq artifact: {error}"),
Self::Quantization(error) => write!(f, "cannot quantize artifact matrix: {error}"),
}
}
}
impl std::error::Error for StreamingConversionError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Source(error) => Some(error),
Self::SourceCensus(report) => Some(report),
Self::Plan(error) => Some(error),
Self::Artifact(error) => Some(error),
Self::Quantization(error) => Some(error),
Self::SourceDigestMismatch { .. } => None,
}
}
}
#[derive(Clone, Debug, PartialEq)]
pub enum QuantizationError {
OutputLength {
values: usize,
output: usize,
},
NonFiniteValue {
index: usize,
value: f32,
},
}
impl fmt::Display for QuantizationError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::OutputLength { values, output } => write!(
f,
"Q8 output length {output} does not match input row length {values}"
),
Self::NonFiniteValue { index, value } => {
write!(
f,
"Q8 input row has non-finite value {value} at index {index}"
)
}
}
}
}
impl std::error::Error for QuantizationError {}
pub trait Q8RowSink {
type Error;
fn write_q8_row(&mut self, row: usize, scale: f32, values: &[i8]) -> Result<(), Self::Error>;
}
#[derive(Clone, Debug, PartialEq)]
pub enum MatrixQuantizationError<E> {
ExpectedMatrix {
rank: usize,
},
EmptyOutputChannel {
shape: Vec<usize>,
},
OutputChannelTooWide {
width: usize,
limit: usize,
},
SourceRowUnavailable {
row: usize,
},
Quantization {
row: usize,
source: QuantizationError,
},
Sink {
row: usize,
source: E,
},
}
impl<E: fmt::Display> fmt::Display for MatrixQuantizationError<E> {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::ExpectedMatrix { rank } => {
write!(
f,
"Q8 matrix quantization requires rank 2 or greater, got rank {rank}"
)
}
Self::EmptyOutputChannel { shape } => write!(
f,
"Q8 matrix quantization refuses empty output channels for shape {shape:?}"
),
Self::OutputChannelTooWide { width, limit } => write!(
f,
"Q8 output-channel width {width} exceeds the bounded adapter limit {limit}"
),
Self::SourceRowUnavailable { row } => {
write!(f, "Q8 source row {row} is unavailable or incomplete")
}
Self::Quantization { row, source } => {
write!(f, "Q8 source row {row} cannot be quantized: {source}")
}
Self::Sink { row, source } => {
write!(f, "Q8 destination rejected row {row}: {source}")
}
}
}
}
impl<E> std::error::Error for MatrixQuantizationError<E>
where
E: std::error::Error + 'static,
{
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Quantization { source, .. } => Some(source),
Self::Sink { source, .. } => Some(source),
Self::ExpectedMatrix { .. }
| Self::EmptyOutputChannel { .. }
| Self::OutputChannelTooWide { .. }
| Self::SourceRowUnavailable { .. } => None,
}
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
pub enum Q8SectionSinkError {
OutputChannelCountTooLarge {
rows: usize,
limit: usize,
},
OutputChannelTooWide {
width: usize,
limit: usize,
},
RowOutOfOrder {
expected: usize,
actual: usize,
},
Incomplete {
expected: usize,
written: usize,
},
Artifact(FttsqError),
}
impl fmt::Display for Q8SectionSinkError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
match self {
Self::OutputChannelCountTooLarge { rows, limit } => write!(
f,
"Q8 matrix has {rows} output channels, exceeding the bounded scale-tail limit {limit}"
),
Self::OutputChannelTooWide { width, limit } => write!(
f,
"Q8 section row width {width} exceeds the bounded row limit {limit}"
),
Self::RowOutOfOrder { expected, actual } => write!(
f,
"Q8 section expected source row {expected}, received row {actual}"
),
Self::Incomplete { expected, written } => write!(
f,
"Q8 section needs {expected} scales but received {written}"
),
Self::Artifact(error) => write!(f, "cannot write Q8 section: {error}"),
}
}
}
impl std::error::Error for Q8SectionSinkError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Artifact(error) => Some(error),
Self::OutputChannelCountTooLarge { .. }
| Self::OutputChannelTooWide { .. }
| Self::RowOutOfOrder { .. }
| Self::Incomplete { .. } => None,
}
}
}
pub struct Q8SectionSink<'a, W> {
writer: &'a mut FttsqStreamingWriter<W>,
section: String,
expected_rows: usize,
next_row: usize,
value_bytes: Vec<u8>,
scale_bytes: Vec<u8>,
}
impl<'a, W: std::io::Write + std::io::Seek> Q8SectionSink<'a, W> {
pub fn new(
writer: &'a mut FttsqStreamingWriter<W>,
section: impl Into<String>,
expected_rows: usize,
) -> Result<Self, Q8SectionSinkError> {
if expected_rows > MAX_Q8_OUTPUT_CHANNELS {
return Err(Q8SectionSinkError::OutputChannelCountTooLarge {
rows: expected_rows,
limit: MAX_Q8_OUTPUT_CHANNELS,
});
}
Ok(Self {
writer,
section: section.into(),
expected_rows,
next_row: 0,
value_bytes: Vec::new(),
scale_bytes: Vec::with_capacity(expected_rows * std::mem::size_of::<f32>()),
})
}
pub fn finish(self) -> Result<(), Q8SectionSinkError> {
if self.next_row != self.expected_rows {
return Err(Q8SectionSinkError::Incomplete {
expected: self.expected_rows,
written: self.next_row,
});
}
self.writer
.write_section(&self.section, &self.scale_bytes)
.map_err(Q8SectionSinkError::Artifact)
}
}
impl<W: std::io::Write + std::io::Seek> Q8RowSink for Q8SectionSink<'_, W> {
type Error = Q8SectionSinkError;
fn write_q8_row(&mut self, row: usize, scale: f32, values: &[i8]) -> Result<(), Self::Error> {
if row != self.next_row {
return Err(Q8SectionSinkError::RowOutOfOrder {
expected: self.next_row,
actual: row,
});
}
if values.len() > MAX_Q8_OUTPUT_CHANNEL_WIDTH {
return Err(Q8SectionSinkError::OutputChannelTooWide {
width: values.len(),
limit: MAX_Q8_OUTPUT_CHANNEL_WIDTH,
});
}
self.value_bytes.clear();
self.value_bytes.extend(
values
.iter()
.map(|&value| u8::from_ne_bytes(value.to_ne_bytes())),
);
self.writer
.write_section(&self.section, &self.value_bytes)
.map_err(Q8SectionSinkError::Artifact)?;
self.scale_bytes.extend_from_slice(&scale.to_le_bytes());
self.next_row += 1;
Ok(())
}
}
pub fn stream_matrix_q8_section<W: std::io::Write + std::io::Seek>(
matrix: &TensorView<'_>,
writer: &mut FttsqStreamingWriter<W>,
section: &str,
) -> Result<(), MatrixQuantizationError<Q8SectionSinkError>> {
let shape = matrix.shape();
if shape.len() < 2 {
return Err(MatrixQuantizationError::ExpectedMatrix { rank: shape.len() });
}
let Some(&row_count) = shape.first() else {
return Err(MatrixQuantizationError::ExpectedMatrix { rank: 0 });
};
let mut sink = Q8SectionSink::new(writer, section, row_count)
.map_err(|source| MatrixQuantizationError::Sink { row: 0, source })?;
quantize_matrix_q8_rows(matrix, &mut sink)?;
sink.finish()
.map_err(|source| MatrixQuantizationError::Sink {
row: row_count,
source,
})
}
pub fn convert_safetensors_streaming<W: std::io::Write + std::io::Seek>(
source: &[u8],
manifest: &WeightsManifest,
plan: &StreamingConversionPlan,
destination: W,
) -> Result<W, StreamingConversionError> {
let index = SafetensorsIndex::parse(source).map_err(StreamingConversionError::Source)?;
manifest
.verify(&index)
.map_err(StreamingConversionError::SourceCensus)?;
let actual_digest = sha256_hex(source);
if actual_digest != plan.source_sha256 {
return Err(StreamingConversionError::SourceDigestMismatch {
expected: plan.source_sha256.clone(),
actual: actual_digest,
});
}
let artifact_plan =
build_artifact_plan(&index, plan).map_err(StreamingConversionError::Plan)?;
let mut writer = artifact_plan
.begin(destination)
.map_err(StreamingConversionError::Artifact)?;
for tensor in tensors_in_write_order(plan) {
let matrix_or_values = index.view(&tensor.source_name, source).ok_or_else(|| {
StreamingConversionError::Plan(ConversionPlanError::SourceTensorMissing {
name: tensor.source_name.clone(),
})
})?;
let section = tensor.section_name();
match tensor.storage {
TensorStoragePolicy::Verbatim => writer
.write_section(section, matrix_or_values.as_bytes())
.map_err(StreamingConversionError::Artifact)?,
TensorStoragePolicy::Q8PerOutputChannel => {
stream_matrix_q8_section(&matrix_or_values, &mut writer, section)
.map_err(StreamingConversionError::Quantization)?;
}
}
}
writer.finish().map_err(StreamingConversionError::Artifact)
}
fn build_artifact_plan(
index: &SafetensorsIndex,
plan: &StreamingConversionPlan,
) -> Result<FttsqStreamPlan, ConversionPlanError> {
if plan.tensors.is_empty() {
return Err(ConversionPlanError::NoTensorPolicies);
}
let mut seen_sources = BTreeSet::<String>::new();
let mut seen_artifacts = BTreeSet::<String>::new();
for tensor in &plan.tensors {
if !seen_sources.insert(tensor.source_name.clone()) {
return Err(ConversionPlanError::DuplicateSourcePolicy {
name: tensor.source_name.clone(),
});
}
if index.entry(&tensor.source_name).is_none() {
return Err(ConversionPlanError::SourceTensorMissing {
name: tensor.source_name.clone(),
});
}
if tensor.artifact_name.is_empty() {
return Err(ConversionPlanError::EmptyArtifactTensorName {
source_name: tensor.source_name.clone(),
});
}
if !seen_artifacts.insert(tensor.artifact_name.clone()) {
return Err(ConversionPlanError::DuplicateArtifactTensor {
name: tensor.artifact_name.clone(),
});
}
if tensor.storage == TensorStoragePolicy::Q8PerOutputChannel {
let entry = index.entry(&tensor.source_name).ok_or_else(|| {
ConversionPlanError::SourceTensorMissing {
name: tensor.source_name.clone(),
}
})?;
if entry.shape.len() < 2 {
return Err(ConversionPlanError::Q8RequiresMatrix {
name: tensor.source_name.clone(),
rank: entry.shape.len(),
});
}
let Some((&rows, trailing_shape)) = entry.shape.split_first() else {
return Err(ConversionPlanError::Q8RequiresMatrix {
name: tensor.source_name.clone(),
rank: 0,
});
};
let row_width = trailing_shape
.iter()
.try_fold(1_usize, |product, &dimension| {
product.checked_mul(dimension)
})
.ok_or_else(|| ConversionPlanError::ShapeOutOfRange {
name: tensor.source_name.clone(),
})?;
if row_width == 0 {
return Err(ConversionPlanError::Q8EmptyOutputChannel {
name: tensor.source_name.clone(),
});
}
if row_width > MAX_Q8_OUTPUT_CHANNEL_WIDTH {
return Err(ConversionPlanError::Q8OutputChannelTooWide {
name: tensor.source_name.clone(),
width: row_width,
limit: MAX_Q8_OUTPUT_CHANNEL_WIDTH,
});
}
if rows > MAX_Q8_OUTPUT_CHANNELS {
return Err(ConversionPlanError::Q8OutputChannelCountTooLarge {
name: tensor.source_name.clone(),
rows,
limit: MAX_Q8_OUTPUT_CHANNELS,
});
}
let scales_name = tensor.scales_name();
if !seen_artifacts.insert(scales_name.clone()) {
return Err(ConversionPlanError::DuplicateArtifactTensor { name: scales_name });
}
}
}
for entry in index.entries() {
if !seen_sources.contains(&entry.name) {
return Err(ConversionPlanError::SourceTensorUnplanned {
name: entry.name.clone(),
});
}
}
let mut artifact_plan = FttsqStreamPlan::new(&plan.model_family, &plan.source_sha256)
.license_notice(&plan.license_notice)
.model_config(plan.model_config.clone())
.quantization_manifest(plan.quantization_manifest.clone());
let mut section_offsets: std::collections::BTreeMap<&'static str, u64> =
std::collections::BTreeMap::new();
let mut declared_sections: Vec<&'static str> = Vec::new();
for tensor in tensors_in_write_order(plan) {
let entry = index.entry(&tensor.source_name).ok_or_else(|| {
ConversionPlanError::SourceTensorMissing {
name: tensor.source_name.clone(),
}
})?;
let shape = artifact_shape(entry, &tensor.source_name)?;
let section = tensor.section_name();
if !declared_sections.contains(§ion) {
declared_sections.push(section);
}
let running = section_offsets.entry(section).or_insert(0);
match tensor.storage {
TensorStoragePolicy::Verbatim => {
let length = u64::try_from(entry.byte_len()).map_err(|_| {
ConversionPlanError::SectionLengthOverflow {
name: tensor.source_name.clone(),
}
})?;
artifact_plan = artifact_plan.tensor(ArtifactTensorEntry {
name: tensor.artifact_name.clone(),
section: section.to_owned(),
dtype: stored_dtype(entry.dtype),
shape,
offset: *running,
length,
scales: None,
});
*running = running.checked_add(length).ok_or_else(|| {
ConversionPlanError::SectionLengthOverflow {
name: tensor.source_name.clone(),
}
})?;
}
TensorStoragePolicy::Q8PerOutputChannel => {
let rows = entry.shape.first().copied().ok_or_else(|| {
ConversionPlanError::Q8RequiresMatrix {
name: tensor.source_name.clone(),
rank: entry.shape.len(),
}
})?;
let values_len = u64::try_from(entry.element_count()).map_err(|_| {
ConversionPlanError::SectionLengthOverflow {
name: tensor.source_name.clone(),
}
})?;
let scales_len = u64::try_from(rows)
.ok()
.and_then(|rows| {
rows.checked_mul(u64::try_from(std::mem::size_of::<f32>()).ok()?)
})
.ok_or_else(|| ConversionPlanError::SectionLengthOverflow {
name: tensor.source_name.clone(),
})?;
let section_len = values_len.checked_add(scales_len).ok_or_else(|| {
ConversionPlanError::SectionLengthOverflow {
name: tensor.source_name.clone(),
}
})?;
let scales_name = tensor.scales_name();
artifact_plan = artifact_plan
.tensor(ArtifactTensorEntry {
name: tensor.artifact_name.clone(),
section: section.to_owned(),
dtype: StoredDtype::Q8,
shape,
offset: *running,
length: values_len,
scales: Some(scales_name.clone()),
})
.tensor(ArtifactTensorEntry {
name: scales_name,
section: section.to_owned(),
dtype: StoredDtype::F32,
shape: vec![u64::try_from(rows).map_err(|_| {
ConversionPlanError::ShapeOutOfRange {
name: tensor.source_name.clone(),
}
})?],
offset: running.checked_add(values_len).ok_or_else(|| {
ConversionPlanError::SectionLengthOverflow {
name: tensor.source_name.clone(),
}
})?,
length: scales_len,
scales: None,
});
*running = running.checked_add(section_len).ok_or_else(|| {
ConversionPlanError::SectionLengthOverflow {
name: tensor.source_name.clone(),
}
})?;
}
}
}
for section in declared_sections {
let class = section_access_class(section);
let length = section_offsets
.get(section)
.copied()
.expect("declared sections accumulate a length");
artifact_plan = artifact_plan.section(section, class, length);
}
Ok(artifact_plan)
}
fn tensors_in_write_order(plan: &StreamingConversionPlan) -> Vec<&TensorConversion> {
let mut order: Vec<&'static str> = Vec::new();
for tensor in &plan.tensors {
let section = tensor.section_name();
if !order.contains(§ion) {
order.push(section);
}
}
let mut grouped = Vec::with_capacity(plan.tensors.len());
for section in order {
grouped.extend(
plan.tensors
.iter()
.filter(|tensor| tensor.section_name() == section),
);
}
grouped
}
fn section_access_class(name: &str) -> AccessClass {
for class in [
AccessClass::HotRecurrentMicrodecoder,
AccessClass::HotRecurrentTalker,
AccessClass::HotCodecDecoder,
AccessClass::ColdTextEmbedding,
AccessClass::EnrollmentSpeakerEncoder,
AccessClass::EnrollmentCodecEncoder,
AccessClass::Metadata,
] {
if class.as_str() == name {
return class;
}
}
unreachable!("section names are minted from AccessClass::as_str")
}
fn artifact_shape(
entry: &crate::safetensors::TensorEntry,
source_name: &str,
) -> Result<Vec<u64>, ConversionPlanError> {
entry
.shape
.iter()
.copied()
.map(u64::try_from)
.collect::<Result<Vec<_>, _>>()
.map_err(|_| ConversionPlanError::ShapeOutOfRange {
name: source_name.to_owned(),
})
}
const fn stored_dtype(source: Dtype) -> StoredDtype {
match source {
Dtype::Bf16 => StoredDtype::Bf16,
Dtype::F32 => StoredDtype::F32,
}
}
fn sha256_hex(bytes: &[u8]) -> String {
const HEX: &[u8; 16] = b"0123456789abcdef";
let digest = {
let mut hasher = Sha256::new();
hasher.update(bytes);
hasher.finish()
};
let mut output = String::with_capacity(64);
for byte in digest {
output.push(char::from(HEX[usize::from(byte >> 4)]));
output.push(char::from(HEX[usize::from(byte & 0x0f)]));
}
output
}
pub fn quantize_output_channel_q8(
row: &[f32],
output: &mut [i8],
) -> Result<f32, QuantizationError> {
if output.len() != row.len() {
return Err(QuantizationError::OutputLength {
values: row.len(),
output: output.len(),
});
}
let mut maximum = 0.0_f32;
for (index, &value) in row.iter().enumerate() {
if !value.is_finite() {
return Err(QuantizationError::NonFiniteValue { index, value });
}
maximum = maximum.max(value.abs());
}
if maximum == 0.0 {
output.fill(0);
return Ok(1.0);
}
let scale = maximum / 127.0;
for (&value, slot) in row.iter().zip(output) {
let rounded = (value / scale).clamp(-127.0, 127.0).round_ties_even();
*slot = rounded as i8;
}
Ok(scale)
}
pub fn quantize_matrix_q8_rows<S: Q8RowSink>(
matrix: &TensorView<'_>,
sink: &mut S,
) -> Result<(), MatrixQuantizationError<S::Error>> {
let shape = matrix.shape();
if shape.len() < 2 {
return Err(MatrixQuantizationError::ExpectedMatrix { rank: shape.len() });
}
let Some(&row_count) = shape.first() else {
return Err(MatrixQuantizationError::ExpectedMatrix { rank: 0 });
};
let row_width = matrix.row_len();
if row_width == 0 {
return Err(MatrixQuantizationError::EmptyOutputChannel {
shape: shape.to_vec(),
});
}
if row_width > MAX_Q8_OUTPUT_CHANNEL_WIDTH {
return Err(MatrixQuantizationError::OutputChannelTooWide {
width: row_width,
limit: MAX_Q8_OUTPUT_CHANNEL_WIDTH,
});
}
let mut source_row = vec![0.0_f32; row_width];
let mut quantized_row = vec![0_i8; row_width];
for row in 0..row_count {
if !matrix.copy_row_f32(row, &mut source_row) {
return Err(MatrixQuantizationError::SourceRowUnavailable { row });
}
let scale = quantize_output_channel_q8(&source_row, &mut quantized_row)
.map_err(|source| MatrixQuantizationError::Quantization { row, source })?;
sink.write_q8_row(row, scale, &quantized_row)
.map_err(|source| MatrixQuantizationError::Sink { row, source })?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use super::*;
use crate::census::ExpectedTensor;
use crate::fttsq::{AccessClass, FttsqReader, FttsqStreamPlan, StoredDtype, TensorEntry};
use crate::safetensors::SafetensorsIndex;
use serde_json::json;
use std::convert::Infallible;
use std::io::Cursor;
#[derive(Default)]
struct RecordingSink {
rows: Vec<(usize, f32, Vec<i8>)>,
}
impl Q8RowSink for RecordingSink {
type Error = Infallible;
fn write_q8_row(
&mut self,
row: usize,
scale: f32,
values: &[i8],
) -> Result<(), Self::Error> {
self.rows.push((row, scale, values.to_vec()));
Ok(())
}
}
fn f32_matrix(rows: usize, columns: usize, values: &[f32]) -> Vec<u8> {
assert_eq!(values.len(), rows * columns);
let payload: Vec<u8> = values
.iter()
.flat_map(|value| value.to_le_bytes())
.collect();
let header = serde_json::to_vec(&json!({
"matrix": {
"dtype": "F32",
"shape": [rows, columns],
"data_offsets": [0, payload.len()],
}
}))
.expect("fixture directory serializes");
let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
bytes.extend_from_slice(&header);
bytes.extend_from_slice(&payload);
bytes
}
fn safetensors(parts: &[(&str, Dtype, &[usize], &[u8])]) -> Vec<u8> {
let mut directory = serde_json::Map::new();
let mut payload = Vec::new();
for (name, dtype, shape, bytes) in parts {
let begin = payload.len();
payload.extend_from_slice(bytes);
directory.insert(
(*name).to_owned(),
json!({
"dtype": dtype.as_str(),
"shape": shape,
"data_offsets": [begin, payload.len()],
}),
);
}
let header = serde_json::to_vec(&serde_json::Value::Object(directory))
.expect("fixture directory serializes");
let mut source = (header.len() as u64).to_le_bytes().to_vec();
source.extend_from_slice(&header);
source.extend_from_slice(&payload);
source
}
#[test]
fn q8_uses_symmetric_ties_to_even_rounding_and_never_emits_negative_128() {
let row = [
-127.0, -126.5, -125.5, -1.5, -0.5, 0.5, 1.5, 125.5, 126.5, 127.0,
];
let mut output = [0_i8; 10];
let scale = quantize_output_channel_q8(&row, &mut output).expect("finite row");
assert_eq!(scale, 1.0);
assert_eq!(output, [-127, -126, -126, -2, 0, 0, 2, 126, 126, 127]);
assert!(!output.contains(&i8::MIN));
}
#[test]
fn q8_all_zero_row_has_a_finite_unit_scale() {
let row = [0.0_f32; 4];
let mut output = [9_i8; 4];
let scale = quantize_output_channel_q8(&row, &mut output).expect("zero row is valid");
assert_eq!(scale, 1.0);
assert_eq!(output, [0; 4]);
}
#[test]
fn q8_refuses_length_mismatch_and_non_finite_input() {
let error = quantize_output_channel_q8(&[1.0, 2.0], &mut [0]).expect_err("wrong length");
assert_eq!(
error,
QuantizationError::OutputLength {
values: 2,
output: 1,
}
);
let error = quantize_output_channel_q8(&[1.0, f32::NAN], &mut [0; 2])
.expect_err("NaN cannot be quantized deterministically");
assert!(matches!(
error,
QuantizationError::NonFiniteValue { index: 1, value } if value.is_nan()
));
}
#[test]
fn runtime_and_offline_callers_receive_byte_identical_q8_rows() {
let row = [-3.0_f32, -0.75, 0.5, 1.5, 3.0];
let mut runtime = [0_i8; 5];
let mut offline = [0_i8; 5];
let runtime_scale = quantize_output_channel_q8(&row, &mut runtime).expect("runtime Q8");
let offline_scale = quantize_output_channel_q8(&row, &mut offline).expect("offline Q8");
assert_eq!(runtime, offline);
assert_eq!(runtime_scale.to_bits(), offline_scale.to_bits());
}
#[test]
fn matrix_rows_stream_through_the_shared_primitive_in_order() {
let bytes = f32_matrix(2, 3, &[1.0, -2.0, 0.5, 3.0, 0.0, -3.0]);
let index = SafetensorsIndex::parse(&bytes).expect("fixture parses");
let matrix = index.view("matrix", &bytes).expect("matrix view exists");
let mut sink = RecordingSink::default();
quantize_matrix_q8_rows(&matrix, &mut sink).expect("finite matrix quantizes");
assert_eq!(sink.rows.len(), 2);
assert_eq!(sink.rows[0].0, 0);
assert_eq!(sink.rows[0].1.to_bits(), (2.0_f32 / 127.0).to_bits());
assert_eq!(sink.rows[0].2, vec![64, -127, 32]);
assert_eq!(sink.rows[1].0, 1);
assert_eq!(sink.rows[1].1.to_bits(), (3.0_f32 / 127.0).to_bits());
assert_eq!(sink.rows[1].2, vec![127, 0, -127]);
}
#[test]
fn matrix_q8_section_streams_values_then_bounded_scale_tail() {
let source = f32_matrix(2, 3, &[1.0, -2.0, 0.5, 3.0, 0.0, -3.0]);
let index = SafetensorsIndex::parse(&source).expect("fixture parses");
let matrix = index.view("matrix", &source).expect("matrix view exists");
let plan = FttsqStreamPlan::new("test-model", "a".repeat(64))
.license_notice("Copyright 2026 Alibaba Cloud\nApache-2.0")
.section("matrix", AccessClass::HotRecurrentTalker, 14)
.tensor(TensorEntry {
name: "matrix.weight".to_owned(),
section: "matrix".to_owned(),
dtype: StoredDtype::Q8,
shape: vec![2, 3],
offset: 0,
length: 6,
scales: Some("matrix.weight.scales".to_owned()),
})
.tensor(TensorEntry {
name: "matrix.weight.scales".to_owned(),
section: "matrix".to_owned(),
dtype: StoredDtype::F32,
shape: vec![2],
offset: 6,
length: 8,
scales: None,
});
let mut writer = plan
.begin(Cursor::new(Vec::new()))
.expect("section metadata is valid");
stream_matrix_q8_section(&matrix, &mut writer, "matrix")
.expect("matrix streams through the canonical Q8 primitive");
let artifact = writer
.finish()
.expect("completed section finalizes its digest")
.into_inner();
let reader = FttsqReader::open(&artifact).expect("artifact verifies");
assert_eq!(
reader
.tensor_bytes("matrix.weight", &artifact)
.expect("Q8 bytes resolve"),
&[64, 129, 32, 127, 0, 129]
);
let scales = reader
.tensor_bytes("matrix.weight.scales", &artifact)
.expect("scale bytes resolve");
assert_eq!(
scales,
&[
(2.0_f32 / 127.0).to_le_bytes(),
(3.0_f32 / 127.0).to_le_bytes(),
]
.concat()
);
}
#[test]
fn manifest_verified_multi_tensor_stream_is_deterministic_and_verbatim_where_required() {
let weight = [1.0_f32, -2.0, 0.5, 3.0, 0.0, -3.0]
.iter()
.flat_map(|value| value.to_le_bytes())
.collect::<Vec<_>>();
let bias = [0x80_u16, 0x3f80]
.iter()
.flat_map(|value| value.to_le_bytes())
.collect::<Vec<_>>();
let source = safetensors(&[
("weight", Dtype::F32, &[2, 3], &weight),
("bias", Dtype::Bf16, &[2], &bias),
]);
let manifest = WeightsManifest::from_expectations(
"small pinned fixture",
[
ExpectedTensor::new("weight", vec![2, 3], Dtype::F32),
ExpectedTensor::new("bias", vec![2], Dtype::Bf16),
],
);
let plan = StreamingConversionPlan::new("qwen3-tts-fixture", sha256_hex(&source))
.license_notice("Copyright 2026 Alibaba Cloud\nApache-2.0")
.model_config(json!({ "fixture": true }))
.quantization_manifest(json!({
"weight": "q8_per_output_channel",
"bias": "verbatim_bf16",
}))
.tensor(TensorConversion::q8_per_output_channel(
"weight",
"weight",
AccessClass::HotRecurrentTalker,
))
.tensor(TensorConversion::verbatim(
"bias",
"bias",
AccessClass::Metadata,
));
let first =
convert_safetensors_streaming(&source, &manifest, &plan, Cursor::new(Vec::new()))
.expect("fixture converts")
.into_inner();
let second =
convert_safetensors_streaming(&source, &manifest, &plan, Cursor::new(Vec::new()))
.expect("second fixture conversion is deterministic")
.into_inner();
assert_eq!(
first, second,
"identical source and plan must be byte-identical"
);
let reader = FttsqReader::open(&first).expect("artifact verifies its section digests");
let mut runtime_q8 = [0_i8; 6];
let runtime_first_scale =
quantize_output_channel_q8(&[1.0_f32, -2.0, 0.5], &mut runtime_q8[..3])
.expect("shared runtime primitive quantizes the first row");
let runtime_second_scale =
quantize_output_channel_q8(&[3.0_f32, 0.0, -3.0], &mut runtime_q8[3..])
.expect("shared runtime primitive quantizes the second row");
assert_eq!(
reader
.tensor_bytes("weight", &first)
.expect("Q8 weights resolve"),
runtime_q8.map(|value| value as u8)
);
assert_eq!(
reader
.tensor_bytes("weight.scales", &first)
.expect("Q8 scales resolve"),
&[
runtime_first_scale.to_le_bytes(),
runtime_second_scale.to_le_bytes(),
]
.concat()
);
assert_eq!(
reader
.tensor_bytes("bias", &first)
.expect("protected BF16 values resolve"),
bias
);
}
#[test]
fn streaming_conversion_refuses_unpinned_source_before_writing() {
let source = f32_matrix(1, 2, &[1.0, -1.0]);
let manifest = WeightsManifest::from_expectations(
"digest fixture",
[ExpectedTensor::new("matrix", vec![1, 2], Dtype::F32)],
);
let plan = StreamingConversionPlan::new("qwen3-tts-fixture", "0".repeat(64))
.license_notice("Copyright 2026 Alibaba Cloud\nApache-2.0")
.tensor(TensorConversion::q8_per_output_channel(
"matrix",
"matrix",
AccessClass::HotRecurrentTalker,
));
let error =
convert_safetensors_streaming(&source, &manifest, &plan, Cursor::new(Vec::new()))
.expect_err("a wrong source digest must refuse before artifact construction");
assert!(matches!(
error,
StreamingConversionError::SourceDigestMismatch { .. }
));
}
#[test]
fn matrix_quantization_refuses_vector_policy_ambiguity() {
let header = serde_json::to_vec(&json!({
"vector": {
"dtype": "F32",
"shape": [2],
"data_offsets": [0, 8],
}
}))
.expect("fixture directory serializes");
let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
bytes.extend_from_slice(&header);
bytes.extend_from_slice(&1.0_f32.to_le_bytes());
bytes.extend_from_slice(&2.0_f32.to_le_bytes());
let index = SafetensorsIndex::parse(&bytes).expect("fixture parses");
let vector = index.view("vector", &bytes).expect("vector view exists");
let error = quantize_matrix_q8_rows(&vector, &mut RecordingSink::default())
.expect_err("vector policy must be explicit");
assert_eq!(error, MatrixQuantizationError::ExpectedMatrix { rank: 1 });
}
#[test]
fn matrix_quantization_refuses_a_row_that_breaks_its_memory_ceiling() {
let values = vec![0.0_f32; MAX_Q8_OUTPUT_CHANNEL_WIDTH + 1];
let bytes = f32_matrix(1, values.len(), &values);
let index = SafetensorsIndex::parse(&bytes).expect("fixture parses");
let matrix = index.view("matrix", &bytes).expect("matrix view exists");
let error = quantize_matrix_q8_rows(&matrix, &mut RecordingSink::default())
.expect_err("row width must be bounded before scratch allocation");
assert_eq!(
error,
MatrixQuantizationError::OutputChannelTooWide {
width: MAX_Q8_OUTPUT_CHANNEL_WIDTH + 1,
limit: MAX_Q8_OUTPUT_CHANNEL_WIDTH,
}
);
}
}