1use std::collections::BTreeSet;
7use std::fmt;
8
9use serde_json::Value;
10
11use crate::census::{CensusReport, WeightsManifest};
12use crate::fttsq::{
13 AccessClass, FttsqError, FttsqStreamPlan, FttsqStreamingWriter, StoredDtype,
14 TensorEntry as ArtifactTensorEntry,
15};
16use crate::safetensors::{Dtype, SafetensorsIndex, TensorView, WeightsError};
17use crate::sha256::Sha256;
18
19pub const MAX_Q8_OUTPUT_CHANNEL_WIDTH: usize = 65_536;
26
27pub const MAX_Q8_OUTPUT_CHANNELS: usize = 262_144;
34
35#[derive(Clone, Copy, Debug, PartialEq, Eq)]
41pub enum TensorStoragePolicy {
42 Verbatim,
44 Q8PerOutputChannel,
46}
47
48#[derive(Clone, Debug, PartialEq, Eq)]
50pub struct TensorConversion {
51 source_name: String,
52 artifact_name: String,
53 access_class: AccessClass,
54 storage: TensorStoragePolicy,
55}
56
57impl TensorConversion {
58 #[must_use]
60 pub fn verbatim(
61 source_name: impl Into<String>,
62 artifact_name: impl Into<String>,
63 access_class: AccessClass,
64 ) -> Self {
65 Self {
66 source_name: source_name.into(),
67 artifact_name: artifact_name.into(),
68 access_class,
69 storage: TensorStoragePolicy::Verbatim,
70 }
71 }
72
73 #[must_use]
75 pub fn q8_per_output_channel(
76 source_name: impl Into<String>,
77 artifact_name: impl Into<String>,
78 access_class: AccessClass,
79 ) -> Self {
80 Self {
81 source_name: source_name.into(),
82 artifact_name: artifact_name.into(),
83 access_class,
84 storage: TensorStoragePolicy::Q8PerOutputChannel,
85 }
86 }
87
88 fn section_name(&self) -> &'static str {
95 self.access_class.as_str()
96 }
97
98 fn scales_name(&self) -> String {
99 format!("{}.scales", self.artifact_name)
100 }
101}
102
103#[derive(Clone, Debug)]
109pub struct StreamingConversionPlan {
110 model_family: String,
111 source_sha256: String,
112 license_notice: String,
113 model_config: Value,
114 quantization_manifest: Value,
115 tensors: Vec<TensorConversion>,
116}
117
118impl StreamingConversionPlan {
119 #[must_use]
121 pub fn new(model_family: impl Into<String>, source_sha256: impl Into<String>) -> Self {
122 Self {
123 model_family: model_family.into(),
124 source_sha256: source_sha256.into(),
125 license_notice: String::new(),
126 model_config: Value::Null,
127 quantization_manifest: Value::Null,
128 tensors: Vec::new(),
129 }
130 }
131
132 #[must_use]
134 pub fn license_notice(mut self, notice: impl Into<String>) -> Self {
135 self.license_notice = notice.into();
136 self
137 }
138
139 #[must_use]
141 pub fn model_config(mut self, config: Value) -> Self {
142 self.model_config = config;
143 self
144 }
145
146 #[must_use]
148 pub fn quantization_manifest(mut self, manifest: Value) -> Self {
149 self.quantization_manifest = manifest;
150 self
151 }
152
153 #[must_use]
155 pub fn tensor(mut self, tensor: TensorConversion) -> Self {
156 self.tensors.push(tensor);
157 self
158 }
159}
160
161#[derive(Clone, Debug, PartialEq, Eq)]
163pub enum ConversionPlanError {
164 NoTensorPolicies,
166 DuplicateSourcePolicy {
168 name: String,
170 },
171 SourceTensorMissing {
173 name: String,
175 },
176 SourceTensorUnplanned {
178 name: String,
180 },
181 DuplicateArtifactTensor {
183 name: String,
185 },
186 EmptyArtifactTensorName {
188 source_name: String,
190 },
191 Q8RequiresMatrix {
193 name: String,
195 rank: usize,
197 },
198 Q8EmptyOutputChannel {
200 name: String,
202 },
203 Q8OutputChannelTooWide {
205 name: String,
207 width: usize,
209 limit: usize,
211 },
212 Q8OutputChannelCountTooLarge {
214 name: String,
216 rows: usize,
218 limit: usize,
220 },
221 ShapeOutOfRange {
223 name: String,
225 },
226 SectionLengthOverflow {
228 name: String,
230 },
231}
232
233impl fmt::Display for ConversionPlanError {
234 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
235 match self {
236 Self::NoTensorPolicies => f.write_str("conversion plan has no tensor policies"),
237 Self::DuplicateSourcePolicy { name } => {
238 write!(
239 f,
240 "conversion plan names source tensor `{name}` more than once"
241 )
242 }
243 Self::SourceTensorMissing { name } => {
244 write!(
245 f,
246 "conversion plan names source tensor `{name}`, which is absent"
247 )
248 }
249 Self::SourceTensorUnplanned { name } => {
250 write!(
251 f,
252 "source tensor `{name}` has no explicit conversion policy"
253 )
254 }
255 Self::DuplicateArtifactTensor { name } => {
256 write!(
257 f,
258 "conversion plan would emit artifact tensor `{name}` more than once"
259 )
260 }
261 Self::EmptyArtifactTensorName { source_name } => write!(
262 f,
263 "conversion plan gives source tensor `{source_name}` an empty artifact name"
264 ),
265 Self::Q8RequiresMatrix { name, rank } => write!(
266 f,
267 "Q8 conversion for `{name}` requires rank 2 or greater, got rank {rank}"
268 ),
269 Self::Q8EmptyOutputChannel { name } => {
270 write!(f, "Q8 conversion for `{name}` has an empty output channel")
271 }
272 Self::Q8OutputChannelTooWide { name, width, limit } => write!(
273 f,
274 "Q8 conversion for `{name}` has output-channel width {width}, exceeding {limit}"
275 ),
276 Self::Q8OutputChannelCountTooLarge { name, rows, limit } => write!(
277 f,
278 "Q8 conversion for `{name}` has {rows} output channels, exceeding {limit}"
279 ),
280 Self::ShapeOutOfRange { name } => {
281 write!(
282 f,
283 "source tensor `{name}` has a shape outside the artifact range"
284 )
285 }
286 Self::SectionLengthOverflow { name } => {
287 write!(
288 f,
289 "source tensor `{name}` overflows its planned artifact section length"
290 )
291 }
292 }
293 }
294}
295
296impl std::error::Error for ConversionPlanError {}
297
298#[derive(Debug)]
300pub enum StreamingConversionError {
301 Source(WeightsError),
303 SourceCensus(Box<CensusReport>),
305 SourceDigestMismatch {
307 expected: String,
309 actual: String,
311 },
312 Plan(ConversionPlanError),
314 Artifact(FttsqError),
316 Quantization(MatrixQuantizationError<Q8SectionSinkError>),
318}
319
320impl fmt::Display for StreamingConversionError {
321 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
322 match self {
323 Self::Source(error) => write!(f, "cannot parse source checkpoint: {error}"),
324 Self::SourceCensus(report) => f.write_str(&report.render()),
325 Self::SourceDigestMismatch { expected, actual } => write!(
326 f,
327 "source checkpoint SHA-256 mismatch: expected {expected}, got {actual}"
328 ),
329 Self::Plan(error) => write!(f, "invalid conversion plan: {error}"),
330 Self::Artifact(error) => write!(f, "cannot write .fttsq artifact: {error}"),
331 Self::Quantization(error) => write!(f, "cannot quantize artifact matrix: {error}"),
332 }
333 }
334}
335
336impl std::error::Error for StreamingConversionError {
337 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
338 match self {
339 Self::Source(error) => Some(error),
340 Self::SourceCensus(report) => Some(report),
341 Self::Plan(error) => Some(error),
342 Self::Artifact(error) => Some(error),
343 Self::Quantization(error) => Some(error),
344 Self::SourceDigestMismatch { .. } => None,
345 }
346 }
347}
348
349#[derive(Clone, Debug, PartialEq)]
351pub enum QuantizationError {
352 OutputLength {
354 values: usize,
356 output: usize,
358 },
359 NonFiniteValue {
361 index: usize,
363 value: f32,
365 },
366}
367
368impl fmt::Display for QuantizationError {
369 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
370 match self {
371 Self::OutputLength { values, output } => write!(
372 f,
373 "Q8 output length {output} does not match input row length {values}"
374 ),
375 Self::NonFiniteValue { index, value } => {
376 write!(
377 f,
378 "Q8 input row has non-finite value {value} at index {index}"
379 )
380 }
381 }
382 }
383}
384
385impl std::error::Error for QuantizationError {}
386
387pub trait Q8RowSink {
394 type Error;
396
397 fn write_q8_row(&mut self, row: usize, scale: f32, values: &[i8]) -> Result<(), Self::Error>;
402}
403
404#[derive(Clone, Debug, PartialEq)]
406pub enum MatrixQuantizationError<E> {
407 ExpectedMatrix {
409 rank: usize,
411 },
412 EmptyOutputChannel {
414 shape: Vec<usize>,
416 },
417 OutputChannelTooWide {
419 width: usize,
421 limit: usize,
423 },
424 SourceRowUnavailable {
429 row: usize,
431 },
432 Quantization {
434 row: usize,
436 source: QuantizationError,
438 },
439 Sink {
441 row: usize,
443 source: E,
445 },
446}
447
448impl<E: fmt::Display> fmt::Display for MatrixQuantizationError<E> {
449 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
450 match self {
451 Self::ExpectedMatrix { rank } => {
452 write!(
453 f,
454 "Q8 matrix quantization requires rank 2 or greater, got rank {rank}"
455 )
456 }
457 Self::EmptyOutputChannel { shape } => write!(
458 f,
459 "Q8 matrix quantization refuses empty output channels for shape {shape:?}"
460 ),
461 Self::OutputChannelTooWide { width, limit } => write!(
462 f,
463 "Q8 output-channel width {width} exceeds the bounded adapter limit {limit}"
464 ),
465 Self::SourceRowUnavailable { row } => {
466 write!(f, "Q8 source row {row} is unavailable or incomplete")
467 }
468 Self::Quantization { row, source } => {
469 write!(f, "Q8 source row {row} cannot be quantized: {source}")
470 }
471 Self::Sink { row, source } => {
472 write!(f, "Q8 destination rejected row {row}: {source}")
473 }
474 }
475 }
476}
477
478impl<E> std::error::Error for MatrixQuantizationError<E>
479where
480 E: std::error::Error + 'static,
481{
482 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
483 match self {
484 Self::Quantization { source, .. } => Some(source),
485 Self::Sink { source, .. } => Some(source),
486 Self::ExpectedMatrix { .. }
487 | Self::EmptyOutputChannel { .. }
488 | Self::OutputChannelTooWide { .. }
489 | Self::SourceRowUnavailable { .. } => None,
490 }
491 }
492}
493
494#[derive(Clone, Debug, PartialEq, Eq)]
496pub enum Q8SectionSinkError {
497 OutputChannelCountTooLarge {
499 rows: usize,
501 limit: usize,
503 },
504 OutputChannelTooWide {
506 width: usize,
508 limit: usize,
510 },
511 RowOutOfOrder {
513 expected: usize,
515 actual: usize,
517 },
518 Incomplete {
520 expected: usize,
522 written: usize,
524 },
525 Artifact(FttsqError),
527}
528
529impl fmt::Display for Q8SectionSinkError {
530 fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
531 match self {
532 Self::OutputChannelCountTooLarge { rows, limit } => write!(
533 f,
534 "Q8 matrix has {rows} output channels, exceeding the bounded scale-tail limit {limit}"
535 ),
536 Self::OutputChannelTooWide { width, limit } => write!(
537 f,
538 "Q8 section row width {width} exceeds the bounded row limit {limit}"
539 ),
540 Self::RowOutOfOrder { expected, actual } => write!(
541 f,
542 "Q8 section expected source row {expected}, received row {actual}"
543 ),
544 Self::Incomplete { expected, written } => write!(
545 f,
546 "Q8 section needs {expected} scales but received {written}"
547 ),
548 Self::Artifact(error) => write!(f, "cannot write Q8 section: {error}"),
549 }
550 }
551}
552
553impl std::error::Error for Q8SectionSinkError {
554 fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
555 match self {
556 Self::Artifact(error) => Some(error),
557 Self::OutputChannelCountTooLarge { .. }
558 | Self::OutputChannelTooWide { .. }
559 | Self::RowOutOfOrder { .. }
560 | Self::Incomplete { .. } => None,
561 }
562 }
563}
564
565pub struct Q8SectionSink<'a, W> {
574 writer: &'a mut FttsqStreamingWriter<W>,
575 section: String,
576 expected_rows: usize,
577 next_row: usize,
578 value_bytes: Vec<u8>,
579 scale_bytes: Vec<u8>,
580}
581
582impl<'a, W: std::io::Write + std::io::Seek> Q8SectionSink<'a, W> {
583 pub fn new(
590 writer: &'a mut FttsqStreamingWriter<W>,
591 section: impl Into<String>,
592 expected_rows: usize,
593 ) -> Result<Self, Q8SectionSinkError> {
594 if expected_rows > MAX_Q8_OUTPUT_CHANNELS {
595 return Err(Q8SectionSinkError::OutputChannelCountTooLarge {
596 rows: expected_rows,
597 limit: MAX_Q8_OUTPUT_CHANNELS,
598 });
599 }
600 Ok(Self {
601 writer,
602 section: section.into(),
603 expected_rows,
604 next_row: 0,
605 value_bytes: Vec::new(),
606 scale_bytes: Vec::with_capacity(expected_rows * std::mem::size_of::<f32>()),
607 })
608 }
609
610 pub fn finish(self) -> Result<(), Q8SectionSinkError> {
616 if self.next_row != self.expected_rows {
617 return Err(Q8SectionSinkError::Incomplete {
618 expected: self.expected_rows,
619 written: self.next_row,
620 });
621 }
622 self.writer
623 .write_section(&self.section, &self.scale_bytes)
624 .map_err(Q8SectionSinkError::Artifact)
625 }
626}
627
628impl<W: std::io::Write + std::io::Seek> Q8RowSink for Q8SectionSink<'_, W> {
629 type Error = Q8SectionSinkError;
630
631 fn write_q8_row(&mut self, row: usize, scale: f32, values: &[i8]) -> Result<(), Self::Error> {
632 if row != self.next_row {
633 return Err(Q8SectionSinkError::RowOutOfOrder {
634 expected: self.next_row,
635 actual: row,
636 });
637 }
638 if values.len() > MAX_Q8_OUTPUT_CHANNEL_WIDTH {
639 return Err(Q8SectionSinkError::OutputChannelTooWide {
640 width: values.len(),
641 limit: MAX_Q8_OUTPUT_CHANNEL_WIDTH,
642 });
643 }
644 self.value_bytes.clear();
645 self.value_bytes.extend(
646 values
647 .iter()
648 .map(|&value| u8::from_ne_bytes(value.to_ne_bytes())),
649 );
650 self.writer
651 .write_section(&self.section, &self.value_bytes)
652 .map_err(Q8SectionSinkError::Artifact)?;
653 self.scale_bytes.extend_from_slice(&scale.to_le_bytes());
654 self.next_row += 1;
655 Ok(())
656 }
657}
658
659pub fn stream_matrix_q8_section<W: std::io::Write + std::io::Seek>(
670 matrix: &TensorView<'_>,
671 writer: &mut FttsqStreamingWriter<W>,
672 section: &str,
673) -> Result<(), MatrixQuantizationError<Q8SectionSinkError>> {
674 let shape = matrix.shape();
675 if shape.len() < 2 {
676 return Err(MatrixQuantizationError::ExpectedMatrix { rank: shape.len() });
677 }
678 let Some(&row_count) = shape.first() else {
679 return Err(MatrixQuantizationError::ExpectedMatrix { rank: 0 });
680 };
681 let mut sink = Q8SectionSink::new(writer, section, row_count)
682 .map_err(|source| MatrixQuantizationError::Sink { row: 0, source })?;
683 quantize_matrix_q8_rows(matrix, &mut sink)?;
684 sink.finish()
685 .map_err(|source| MatrixQuantizationError::Sink {
686 row: row_count,
687 source,
688 })
689}
690
691pub fn convert_safetensors_streaming<W: std::io::Write + std::io::Seek>(
710 source: &[u8],
711 manifest: &WeightsManifest,
712 plan: &StreamingConversionPlan,
713 destination: W,
714) -> Result<W, StreamingConversionError> {
715 let index = SafetensorsIndex::parse(source).map_err(StreamingConversionError::Source)?;
716 manifest
717 .verify(&index)
718 .map_err(StreamingConversionError::SourceCensus)?;
719
720 let actual_digest = sha256_hex(source);
721 if actual_digest != plan.source_sha256 {
722 return Err(StreamingConversionError::SourceDigestMismatch {
723 expected: plan.source_sha256.clone(),
724 actual: actual_digest,
725 });
726 }
727
728 let artifact_plan =
729 build_artifact_plan(&index, plan).map_err(StreamingConversionError::Plan)?;
730 let mut writer = artifact_plan
731 .begin(destination)
732 .map_err(StreamingConversionError::Artifact)?;
733
734 for tensor in tensors_in_write_order(plan) {
735 let matrix_or_values = index.view(&tensor.source_name, source).ok_or_else(|| {
736 StreamingConversionError::Plan(ConversionPlanError::SourceTensorMissing {
737 name: tensor.source_name.clone(),
738 })
739 })?;
740 let section = tensor.section_name();
741 match tensor.storage {
742 TensorStoragePolicy::Verbatim => writer
743 .write_section(section, matrix_or_values.as_bytes())
744 .map_err(StreamingConversionError::Artifact)?,
745 TensorStoragePolicy::Q8PerOutputChannel => {
746 stream_matrix_q8_section(&matrix_or_values, &mut writer, section)
747 .map_err(StreamingConversionError::Quantization)?;
748 }
749 }
750 }
751
752 writer.finish().map_err(StreamingConversionError::Artifact)
753}
754
755fn build_artifact_plan(
756 index: &SafetensorsIndex,
757 plan: &StreamingConversionPlan,
758) -> Result<FttsqStreamPlan, ConversionPlanError> {
759 if plan.tensors.is_empty() {
760 return Err(ConversionPlanError::NoTensorPolicies);
761 }
762
763 let mut seen_sources = BTreeSet::<String>::new();
764 let mut seen_artifacts = BTreeSet::<String>::new();
765 for tensor in &plan.tensors {
766 if !seen_sources.insert(tensor.source_name.clone()) {
767 return Err(ConversionPlanError::DuplicateSourcePolicy {
768 name: tensor.source_name.clone(),
769 });
770 }
771 if index.entry(&tensor.source_name).is_none() {
772 return Err(ConversionPlanError::SourceTensorMissing {
773 name: tensor.source_name.clone(),
774 });
775 }
776 if tensor.artifact_name.is_empty() {
777 return Err(ConversionPlanError::EmptyArtifactTensorName {
778 source_name: tensor.source_name.clone(),
779 });
780 }
781 if !seen_artifacts.insert(tensor.artifact_name.clone()) {
782 return Err(ConversionPlanError::DuplicateArtifactTensor {
783 name: tensor.artifact_name.clone(),
784 });
785 }
786 if tensor.storage == TensorStoragePolicy::Q8PerOutputChannel {
787 let entry = index.entry(&tensor.source_name).ok_or_else(|| {
788 ConversionPlanError::SourceTensorMissing {
789 name: tensor.source_name.clone(),
790 }
791 })?;
792 if entry.shape.len() < 2 {
793 return Err(ConversionPlanError::Q8RequiresMatrix {
794 name: tensor.source_name.clone(),
795 rank: entry.shape.len(),
796 });
797 }
798 let Some((&rows, trailing_shape)) = entry.shape.split_first() else {
799 return Err(ConversionPlanError::Q8RequiresMatrix {
800 name: tensor.source_name.clone(),
801 rank: 0,
802 });
803 };
804 let row_width = trailing_shape
805 .iter()
806 .try_fold(1_usize, |product, &dimension| {
807 product.checked_mul(dimension)
808 })
809 .ok_or_else(|| ConversionPlanError::ShapeOutOfRange {
810 name: tensor.source_name.clone(),
811 })?;
812 if row_width == 0 {
813 return Err(ConversionPlanError::Q8EmptyOutputChannel {
814 name: tensor.source_name.clone(),
815 });
816 }
817 if row_width > MAX_Q8_OUTPUT_CHANNEL_WIDTH {
818 return Err(ConversionPlanError::Q8OutputChannelTooWide {
819 name: tensor.source_name.clone(),
820 width: row_width,
821 limit: MAX_Q8_OUTPUT_CHANNEL_WIDTH,
822 });
823 }
824 if rows > MAX_Q8_OUTPUT_CHANNELS {
825 return Err(ConversionPlanError::Q8OutputChannelCountTooLarge {
826 name: tensor.source_name.clone(),
827 rows,
828 limit: MAX_Q8_OUTPUT_CHANNELS,
829 });
830 }
831 let scales_name = tensor.scales_name();
832 if !seen_artifacts.insert(scales_name.clone()) {
833 return Err(ConversionPlanError::DuplicateArtifactTensor { name: scales_name });
834 }
835 }
836 }
837
838 for entry in index.entries() {
839 if !seen_sources.contains(&entry.name) {
840 return Err(ConversionPlanError::SourceTensorUnplanned {
841 name: entry.name.clone(),
842 });
843 }
844 }
845
846 let mut artifact_plan = FttsqStreamPlan::new(&plan.model_family, &plan.source_sha256)
847 .license_notice(&plan.license_notice)
848 .model_config(plan.model_config.clone())
849 .quantization_manifest(plan.quantization_manifest.clone());
850
851 let mut section_offsets: std::collections::BTreeMap<&'static str, u64> =
854 std::collections::BTreeMap::new();
855 let mut declared_sections: Vec<&'static str> = Vec::new();
856 for tensor in tensors_in_write_order(plan) {
857 let entry = index.entry(&tensor.source_name).ok_or_else(|| {
858 ConversionPlanError::SourceTensorMissing {
859 name: tensor.source_name.clone(),
860 }
861 })?;
862 let shape = artifact_shape(entry, &tensor.source_name)?;
863 let section = tensor.section_name();
864 if !declared_sections.contains(§ion) {
865 declared_sections.push(section);
866 }
867 let running = section_offsets.entry(section).or_insert(0);
868 match tensor.storage {
869 TensorStoragePolicy::Verbatim => {
870 let length = u64::try_from(entry.byte_len()).map_err(|_| {
871 ConversionPlanError::SectionLengthOverflow {
872 name: tensor.source_name.clone(),
873 }
874 })?;
875 artifact_plan = artifact_plan.tensor(ArtifactTensorEntry {
876 name: tensor.artifact_name.clone(),
877 section: section.to_owned(),
878 dtype: stored_dtype(entry.dtype),
879 shape,
880 offset: *running,
881 length,
882 scales: None,
883 });
884 *running = running.checked_add(length).ok_or_else(|| {
885 ConversionPlanError::SectionLengthOverflow {
886 name: tensor.source_name.clone(),
887 }
888 })?;
889 }
890 TensorStoragePolicy::Q8PerOutputChannel => {
891 let rows = entry.shape.first().copied().ok_or_else(|| {
892 ConversionPlanError::Q8RequiresMatrix {
893 name: tensor.source_name.clone(),
894 rank: entry.shape.len(),
895 }
896 })?;
897 let values_len = u64::try_from(entry.element_count()).map_err(|_| {
898 ConversionPlanError::SectionLengthOverflow {
899 name: tensor.source_name.clone(),
900 }
901 })?;
902 let scales_len = u64::try_from(rows)
903 .ok()
904 .and_then(|rows| {
905 rows.checked_mul(u64::try_from(std::mem::size_of::<f32>()).ok()?)
906 })
907 .ok_or_else(|| ConversionPlanError::SectionLengthOverflow {
908 name: tensor.source_name.clone(),
909 })?;
910 let section_len = values_len.checked_add(scales_len).ok_or_else(|| {
911 ConversionPlanError::SectionLengthOverflow {
912 name: tensor.source_name.clone(),
913 }
914 })?;
915 let scales_name = tensor.scales_name();
916 artifact_plan = artifact_plan
917 .tensor(ArtifactTensorEntry {
918 name: tensor.artifact_name.clone(),
919 section: section.to_owned(),
920 dtype: StoredDtype::Q8,
921 shape,
922 offset: *running,
923 length: values_len,
924 scales: Some(scales_name.clone()),
925 })
926 .tensor(ArtifactTensorEntry {
927 name: scales_name,
928 section: section.to_owned(),
929 dtype: StoredDtype::F32,
930 shape: vec![u64::try_from(rows).map_err(|_| {
931 ConversionPlanError::ShapeOutOfRange {
932 name: tensor.source_name.clone(),
933 }
934 })?],
935 offset: running.checked_add(values_len).ok_or_else(|| {
936 ConversionPlanError::SectionLengthOverflow {
937 name: tensor.source_name.clone(),
938 }
939 })?,
940 length: scales_len,
941 scales: None,
942 });
943 *running = running.checked_add(section_len).ok_or_else(|| {
944 ConversionPlanError::SectionLengthOverflow {
945 name: tensor.source_name.clone(),
946 }
947 })?;
948 }
949 }
950 }
951
952 for section in declared_sections {
955 let class = section_access_class(section);
956 let length = section_offsets
957 .get(section)
958 .copied()
959 .expect("declared sections accumulate a length");
960 artifact_plan = artifact_plan.section(section, class, length);
961 }
962
963 Ok(artifact_plan)
964}
965
966fn tensors_in_write_order(plan: &StreamingConversionPlan) -> Vec<&TensorConversion> {
969 let mut order: Vec<&'static str> = Vec::new();
970 for tensor in &plan.tensors {
971 let section = tensor.section_name();
972 if !order.contains(§ion) {
973 order.push(section);
974 }
975 }
976 let mut grouped = Vec::with_capacity(plan.tensors.len());
977 for section in order {
978 grouped.extend(
979 plan.tensors
980 .iter()
981 .filter(|tensor| tensor.section_name() == section),
982 );
983 }
984 grouped
985}
986
987fn section_access_class(name: &str) -> AccessClass {
989 for class in [
990 AccessClass::HotRecurrentMicrodecoder,
991 AccessClass::HotRecurrentTalker,
992 AccessClass::HotCodecDecoder,
993 AccessClass::ColdTextEmbedding,
994 AccessClass::EnrollmentSpeakerEncoder,
995 AccessClass::EnrollmentCodecEncoder,
996 AccessClass::Metadata,
997 ] {
998 if class.as_str() == name {
999 return class;
1000 }
1001 }
1002 unreachable!("section names are minted from AccessClass::as_str")
1003}
1004
1005fn artifact_shape(
1006 entry: &crate::safetensors::TensorEntry,
1007 source_name: &str,
1008) -> Result<Vec<u64>, ConversionPlanError> {
1009 entry
1010 .shape
1011 .iter()
1012 .copied()
1013 .map(u64::try_from)
1014 .collect::<Result<Vec<_>, _>>()
1015 .map_err(|_| ConversionPlanError::ShapeOutOfRange {
1016 name: source_name.to_owned(),
1017 })
1018}
1019
1020const fn stored_dtype(source: Dtype) -> StoredDtype {
1021 match source {
1022 Dtype::Bf16 => StoredDtype::Bf16,
1023 Dtype::F32 => StoredDtype::F32,
1024 }
1025}
1026
1027fn sha256_hex(bytes: &[u8]) -> String {
1028 const HEX: &[u8; 16] = b"0123456789abcdef";
1029 let digest = {
1030 let mut hasher = Sha256::new();
1031 hasher.update(bytes);
1032 hasher.finish()
1033 };
1034 let mut output = String::with_capacity(64);
1035 for byte in digest {
1036 output.push(char::from(HEX[usize::from(byte >> 4)]));
1037 output.push(char::from(HEX[usize::from(byte & 0x0f)]));
1038 }
1039 output
1040}
1041
1042pub fn quantize_output_channel_q8(
1057 row: &[f32],
1058 output: &mut [i8],
1059) -> Result<f32, QuantizationError> {
1060 if output.len() != row.len() {
1061 return Err(QuantizationError::OutputLength {
1062 values: row.len(),
1063 output: output.len(),
1064 });
1065 }
1066
1067 let mut maximum = 0.0_f32;
1068 for (index, &value) in row.iter().enumerate() {
1069 if !value.is_finite() {
1070 return Err(QuantizationError::NonFiniteValue { index, value });
1071 }
1072 maximum = maximum.max(value.abs());
1073 }
1074
1075 if maximum == 0.0 {
1076 output.fill(0);
1077 return Ok(1.0);
1078 }
1079
1080 let scale = maximum / 127.0;
1081 for (&value, slot) in row.iter().zip(output) {
1082 let rounded = (value / scale).clamp(-127.0, 127.0).round_ties_even();
1083 *slot = rounded as i8;
1086 }
1087 Ok(scale)
1088}
1089
1090pub fn quantize_matrix_q8_rows<S: Q8RowSink>(
1107 matrix: &TensorView<'_>,
1108 sink: &mut S,
1109) -> Result<(), MatrixQuantizationError<S::Error>> {
1110 let shape = matrix.shape();
1111 if shape.len() < 2 {
1112 return Err(MatrixQuantizationError::ExpectedMatrix { rank: shape.len() });
1113 }
1114
1115 let Some(&row_count) = shape.first() else {
1116 return Err(MatrixQuantizationError::ExpectedMatrix { rank: 0 });
1117 };
1118 let row_width = matrix.row_len();
1119 if row_width == 0 {
1120 return Err(MatrixQuantizationError::EmptyOutputChannel {
1121 shape: shape.to_vec(),
1122 });
1123 }
1124 if row_width > MAX_Q8_OUTPUT_CHANNEL_WIDTH {
1125 return Err(MatrixQuantizationError::OutputChannelTooWide {
1126 width: row_width,
1127 limit: MAX_Q8_OUTPUT_CHANNEL_WIDTH,
1128 });
1129 }
1130
1131 let mut source_row = vec![0.0_f32; row_width];
1132 let mut quantized_row = vec![0_i8; row_width];
1133 for row in 0..row_count {
1134 if !matrix.copy_row_f32(row, &mut source_row) {
1135 return Err(MatrixQuantizationError::SourceRowUnavailable { row });
1136 }
1137 let scale = quantize_output_channel_q8(&source_row, &mut quantized_row)
1138 .map_err(|source| MatrixQuantizationError::Quantization { row, source })?;
1139 sink.write_q8_row(row, scale, &quantized_row)
1140 .map_err(|source| MatrixQuantizationError::Sink { row, source })?;
1141 }
1142 Ok(())
1143}
1144
1145#[cfg(test)]
1146mod tests {
1147 use super::*;
1148 use crate::census::ExpectedTensor;
1149 use crate::fttsq::{AccessClass, FttsqReader, FttsqStreamPlan, StoredDtype, TensorEntry};
1150 use crate::safetensors::SafetensorsIndex;
1151 use serde_json::json;
1152 use std::convert::Infallible;
1153 use std::io::Cursor;
1154
1155 #[derive(Default)]
1156 struct RecordingSink {
1157 rows: Vec<(usize, f32, Vec<i8>)>,
1158 }
1159
1160 impl Q8RowSink for RecordingSink {
1161 type Error = Infallible;
1162
1163 fn write_q8_row(
1164 &mut self,
1165 row: usize,
1166 scale: f32,
1167 values: &[i8],
1168 ) -> Result<(), Self::Error> {
1169 self.rows.push((row, scale, values.to_vec()));
1170 Ok(())
1171 }
1172 }
1173
1174 fn f32_matrix(rows: usize, columns: usize, values: &[f32]) -> Vec<u8> {
1175 assert_eq!(values.len(), rows * columns);
1176 let payload: Vec<u8> = values
1177 .iter()
1178 .flat_map(|value| value.to_le_bytes())
1179 .collect();
1180 let header = serde_json::to_vec(&json!({
1181 "matrix": {
1182 "dtype": "F32",
1183 "shape": [rows, columns],
1184 "data_offsets": [0, payload.len()],
1185 }
1186 }))
1187 .expect("fixture directory serializes");
1188
1189 let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
1190 bytes.extend_from_slice(&header);
1191 bytes.extend_from_slice(&payload);
1192 bytes
1193 }
1194
1195 fn safetensors(parts: &[(&str, Dtype, &[usize], &[u8])]) -> Vec<u8> {
1196 let mut directory = serde_json::Map::new();
1197 let mut payload = Vec::new();
1198 for (name, dtype, shape, bytes) in parts {
1199 let begin = payload.len();
1200 payload.extend_from_slice(bytes);
1201 directory.insert(
1202 (*name).to_owned(),
1203 json!({
1204 "dtype": dtype.as_str(),
1205 "shape": shape,
1206 "data_offsets": [begin, payload.len()],
1207 }),
1208 );
1209 }
1210 let header = serde_json::to_vec(&serde_json::Value::Object(directory))
1211 .expect("fixture directory serializes");
1212 let mut source = (header.len() as u64).to_le_bytes().to_vec();
1213 source.extend_from_slice(&header);
1214 source.extend_from_slice(&payload);
1215 source
1216 }
1217
1218 #[test]
1219 fn q8_uses_symmetric_ties_to_even_rounding_and_never_emits_negative_128() {
1220 let row = [
1221 -127.0, -126.5, -125.5, -1.5, -0.5, 0.5, 1.5, 125.5, 126.5, 127.0,
1222 ];
1223 let mut output = [0_i8; 10];
1224
1225 let scale = quantize_output_channel_q8(&row, &mut output).expect("finite row");
1226
1227 assert_eq!(scale, 1.0);
1228 assert_eq!(output, [-127, -126, -126, -2, 0, 0, 2, 126, 126, 127]);
1229 assert!(!output.contains(&i8::MIN));
1230 }
1231
1232 #[test]
1233 fn q8_all_zero_row_has_a_finite_unit_scale() {
1234 let row = [0.0_f32; 4];
1235 let mut output = [9_i8; 4];
1236
1237 let scale = quantize_output_channel_q8(&row, &mut output).expect("zero row is valid");
1238
1239 assert_eq!(scale, 1.0);
1240 assert_eq!(output, [0; 4]);
1241 }
1242
1243 #[test]
1244 fn q8_refuses_length_mismatch_and_non_finite_input() {
1245 let error = quantize_output_channel_q8(&[1.0, 2.0], &mut [0]).expect_err("wrong length");
1246 assert_eq!(
1247 error,
1248 QuantizationError::OutputLength {
1249 values: 2,
1250 output: 1,
1251 }
1252 );
1253
1254 let error = quantize_output_channel_q8(&[1.0, f32::NAN], &mut [0; 2])
1255 .expect_err("NaN cannot be quantized deterministically");
1256 assert!(matches!(
1257 error,
1258 QuantizationError::NonFiniteValue { index: 1, value } if value.is_nan()
1259 ));
1260 }
1261
1262 #[test]
1263 fn runtime_and_offline_callers_receive_byte_identical_q8_rows() {
1264 let row = [-3.0_f32, -0.75, 0.5, 1.5, 3.0];
1265 let mut runtime = [0_i8; 5];
1266 let mut offline = [0_i8; 5];
1267
1268 let runtime_scale = quantize_output_channel_q8(&row, &mut runtime).expect("runtime Q8");
1269 let offline_scale = quantize_output_channel_q8(&row, &mut offline).expect("offline Q8");
1270
1271 assert_eq!(runtime, offline);
1272 assert_eq!(runtime_scale.to_bits(), offline_scale.to_bits());
1273 }
1274
1275 #[test]
1276 fn matrix_rows_stream_through_the_shared_primitive_in_order() {
1277 let bytes = f32_matrix(2, 3, &[1.0, -2.0, 0.5, 3.0, 0.0, -3.0]);
1278 let index = SafetensorsIndex::parse(&bytes).expect("fixture parses");
1279 let matrix = index.view("matrix", &bytes).expect("matrix view exists");
1280 let mut sink = RecordingSink::default();
1281
1282 quantize_matrix_q8_rows(&matrix, &mut sink).expect("finite matrix quantizes");
1283
1284 assert_eq!(sink.rows.len(), 2);
1285 assert_eq!(sink.rows[0].0, 0);
1286 assert_eq!(sink.rows[0].1.to_bits(), (2.0_f32 / 127.0).to_bits());
1287 assert_eq!(sink.rows[0].2, vec![64, -127, 32]);
1288 assert_eq!(sink.rows[1].0, 1);
1289 assert_eq!(sink.rows[1].1.to_bits(), (3.0_f32 / 127.0).to_bits());
1290 assert_eq!(sink.rows[1].2, vec![127, 0, -127]);
1291 }
1292
1293 #[test]
1294 fn matrix_q8_section_streams_values_then_bounded_scale_tail() {
1295 let source = f32_matrix(2, 3, &[1.0, -2.0, 0.5, 3.0, 0.0, -3.0]);
1296 let index = SafetensorsIndex::parse(&source).expect("fixture parses");
1297 let matrix = index.view("matrix", &source).expect("matrix view exists");
1298 let plan = FttsqStreamPlan::new("test-model", "a".repeat(64))
1299 .license_notice("Copyright 2026 Alibaba Cloud\nApache-2.0")
1300 .section("matrix", AccessClass::HotRecurrentTalker, 14)
1301 .tensor(TensorEntry {
1302 name: "matrix.weight".to_owned(),
1303 section: "matrix".to_owned(),
1304 dtype: StoredDtype::Q8,
1305 shape: vec![2, 3],
1306 offset: 0,
1307 length: 6,
1308 scales: Some("matrix.weight.scales".to_owned()),
1309 })
1310 .tensor(TensorEntry {
1311 name: "matrix.weight.scales".to_owned(),
1312 section: "matrix".to_owned(),
1313 dtype: StoredDtype::F32,
1314 shape: vec![2],
1315 offset: 6,
1316 length: 8,
1317 scales: None,
1318 });
1319 let mut writer = plan
1320 .begin(Cursor::new(Vec::new()))
1321 .expect("section metadata is valid");
1322
1323 stream_matrix_q8_section(&matrix, &mut writer, "matrix")
1324 .expect("matrix streams through the canonical Q8 primitive");
1325 let artifact = writer
1326 .finish()
1327 .expect("completed section finalizes its digest")
1328 .into_inner();
1329 let reader = FttsqReader::open(&artifact).expect("artifact verifies");
1330
1331 assert_eq!(
1332 reader
1333 .tensor_bytes("matrix.weight", &artifact)
1334 .expect("Q8 bytes resolve"),
1335 &[64, 129, 32, 127, 0, 129]
1336 );
1337 let scales = reader
1338 .tensor_bytes("matrix.weight.scales", &artifact)
1339 .expect("scale bytes resolve");
1340 assert_eq!(
1341 scales,
1342 &[
1343 (2.0_f32 / 127.0).to_le_bytes(),
1344 (3.0_f32 / 127.0).to_le_bytes(),
1345 ]
1346 .concat()
1347 );
1348 }
1349
1350 #[test]
1351 fn manifest_verified_multi_tensor_stream_is_deterministic_and_verbatim_where_required() {
1352 let weight = [1.0_f32, -2.0, 0.5, 3.0, 0.0, -3.0]
1353 .iter()
1354 .flat_map(|value| value.to_le_bytes())
1355 .collect::<Vec<_>>();
1356 let bias = [0x80_u16, 0x3f80]
1357 .iter()
1358 .flat_map(|value| value.to_le_bytes())
1359 .collect::<Vec<_>>();
1360 let source = safetensors(&[
1361 ("weight", Dtype::F32, &[2, 3], &weight),
1362 ("bias", Dtype::Bf16, &[2], &bias),
1363 ]);
1364 let manifest = WeightsManifest::from_expectations(
1365 "small pinned fixture",
1366 [
1367 ExpectedTensor::new("weight", vec![2, 3], Dtype::F32),
1368 ExpectedTensor::new("bias", vec![2], Dtype::Bf16),
1369 ],
1370 );
1371 let plan = StreamingConversionPlan::new("qwen3-tts-fixture", sha256_hex(&source))
1372 .license_notice("Copyright 2026 Alibaba Cloud\nApache-2.0")
1373 .model_config(json!({ "fixture": true }))
1374 .quantization_manifest(json!({
1375 "weight": "q8_per_output_channel",
1376 "bias": "verbatim_bf16",
1377 }))
1378 .tensor(TensorConversion::q8_per_output_channel(
1379 "weight",
1380 "weight",
1381 AccessClass::HotRecurrentTalker,
1382 ))
1383 .tensor(TensorConversion::verbatim(
1384 "bias",
1385 "bias",
1386 AccessClass::Metadata,
1387 ));
1388
1389 let first =
1390 convert_safetensors_streaming(&source, &manifest, &plan, Cursor::new(Vec::new()))
1391 .expect("fixture converts")
1392 .into_inner();
1393 let second =
1394 convert_safetensors_streaming(&source, &manifest, &plan, Cursor::new(Vec::new()))
1395 .expect("second fixture conversion is deterministic")
1396 .into_inner();
1397 assert_eq!(
1398 first, second,
1399 "identical source and plan must be byte-identical"
1400 );
1401
1402 let reader = FttsqReader::open(&first).expect("artifact verifies its section digests");
1403 let mut runtime_q8 = [0_i8; 6];
1404 let runtime_first_scale =
1405 quantize_output_channel_q8(&[1.0_f32, -2.0, 0.5], &mut runtime_q8[..3])
1406 .expect("shared runtime primitive quantizes the first row");
1407 let runtime_second_scale =
1408 quantize_output_channel_q8(&[3.0_f32, 0.0, -3.0], &mut runtime_q8[3..])
1409 .expect("shared runtime primitive quantizes the second row");
1410 assert_eq!(
1411 reader
1412 .tensor_bytes("weight", &first)
1413 .expect("Q8 weights resolve"),
1414 runtime_q8.map(|value| value as u8)
1415 );
1416 assert_eq!(
1417 reader
1418 .tensor_bytes("weight.scales", &first)
1419 .expect("Q8 scales resolve"),
1420 &[
1421 runtime_first_scale.to_le_bytes(),
1422 runtime_second_scale.to_le_bytes(),
1423 ]
1424 .concat()
1425 );
1426 assert_eq!(
1427 reader
1428 .tensor_bytes("bias", &first)
1429 .expect("protected BF16 values resolve"),
1430 bias
1431 );
1432 }
1433
1434 #[test]
1435 fn streaming_conversion_refuses_unpinned_source_before_writing() {
1436 let source = f32_matrix(1, 2, &[1.0, -1.0]);
1437 let manifest = WeightsManifest::from_expectations(
1438 "digest fixture",
1439 [ExpectedTensor::new("matrix", vec![1, 2], Dtype::F32)],
1440 );
1441 let plan = StreamingConversionPlan::new("qwen3-tts-fixture", "0".repeat(64))
1442 .license_notice("Copyright 2026 Alibaba Cloud\nApache-2.0")
1443 .tensor(TensorConversion::q8_per_output_channel(
1444 "matrix",
1445 "matrix",
1446 AccessClass::HotRecurrentTalker,
1447 ));
1448
1449 let error =
1450 convert_safetensors_streaming(&source, &manifest, &plan, Cursor::new(Vec::new()))
1451 .expect_err("a wrong source digest must refuse before artifact construction");
1452 assert!(matches!(
1453 error,
1454 StreamingConversionError::SourceDigestMismatch { .. }
1455 ));
1456 }
1457
1458 #[test]
1459 fn matrix_quantization_refuses_vector_policy_ambiguity() {
1460 let header = serde_json::to_vec(&json!({
1461 "vector": {
1462 "dtype": "F32",
1463 "shape": [2],
1464 "data_offsets": [0, 8],
1465 }
1466 }))
1467 .expect("fixture directory serializes");
1468 let mut bytes = (header.len() as u64).to_le_bytes().to_vec();
1469 bytes.extend_from_slice(&header);
1470 bytes.extend_from_slice(&1.0_f32.to_le_bytes());
1471 bytes.extend_from_slice(&2.0_f32.to_le_bytes());
1472 let index = SafetensorsIndex::parse(&bytes).expect("fixture parses");
1473 let vector = index.view("vector", &bytes).expect("vector view exists");
1474
1475 let error = quantize_matrix_q8_rows(&vector, &mut RecordingSink::default())
1476 .expect_err("vector policy must be explicit");
1477 assert_eq!(error, MatrixQuantizationError::ExpectedMatrix { rank: 1 });
1478 }
1479
1480 #[test]
1481 fn matrix_quantization_refuses_a_row_that_breaks_its_memory_ceiling() {
1482 let values = vec![0.0_f32; MAX_Q8_OUTPUT_CHANNEL_WIDTH + 1];
1483 let bytes = f32_matrix(1, values.len(), &values);
1484 let index = SafetensorsIndex::parse(&bytes).expect("fixture parses");
1485 let matrix = index.view("matrix", &bytes).expect("matrix view exists");
1486
1487 let error = quantize_matrix_q8_rows(&matrix, &mut RecordingSink::default())
1488 .expect_err("row width must be bounded before scratch allocation");
1489 assert_eq!(
1490 error,
1491 MatrixQuantizationError::OutputChannelTooWide {
1492 width: MAX_Q8_OUTPUT_CHANNEL_WIDTH + 1,
1493 limit: MAX_Q8_OUTPUT_CHANNEL_WIDTH,
1494 }
1495 );
1496 }
1497}