1use super::*;
6use crate::container_10d::crc32c::crc32c;
7
8fn decode_safetensor_to_f32(raw: &[u8], ggml: u32, count: usize, out: &mut Vec<f32>) {
11 use crate::safetensor::{GGML_BF16, GGML_F16, GGML_F32};
12 out.clear();
13 out.reserve(count);
14 match ggml {
15 GGML_F32 => {
16 for k in 0..count {
17 let o = k * 4;
18 if o + 4 > raw.len() {
19 break;
20 }
21 out.push(f32::from_le_bytes([
22 raw[o],
23 raw[o + 1],
24 raw[o + 2],
25 raw[o + 3],
26 ]));
27 }
28 }
29 GGML_F16 => {
30 for k in 0..count {
31 let o = k * 2;
32 if o + 2 > raw.len() {
33 break;
34 }
35 out.push(half::f16::from_le_bytes([raw[o], raw[o + 1]]).to_f32());
36 }
37 }
38 GGML_BF16 => {
39 for k in 0..count {
40 let o = k * 2;
41 if o + 2 > raw.len() {
42 break;
43 }
44 out.push(half::bf16::from_le_bytes([raw[o], raw[o + 1]]).to_f32());
45 }
46 }
47 _ => {}
48 }
49}
50
51#[derive(Debug, Clone, Copy, PartialEq, Eq)]
53pub struct TranscodeReport {
54 pub n_tensors: usize,
55 pub bytes_written: usize,
56 pub largest_tensor_bytes: usize,
57 pub total_tensor_bytes: usize,
58 pub peak_working_bytes: usize,
61}
62
63#[derive(Clone, Copy)]
74enum SafetensorQuantization {
75 Verbatim,
76 AllTernary,
77 FfnTernary,
78}
79
80fn transcode_safetensor_with_policy<W: std::io::Write>(
81 src: &[u8],
82 page_log2: u16,
83 out: &mut W,
84 policy: SafetensorQuantization,
85) -> Result<TranscodeReport, String> {
86 use crate::safetensor::{
87 ggml_elem_bytes, is_high_fidelity_ggml, parse_safetensor_header, safetensor_dtype_to_ggml,
88 };
89 use crate::ternary::{ternary_blob, ternary_blob_len, GGML_TYPE_TERNARY_158};
90 let page_log2 = if page_log2 == 0 {
91 P64_DEFAULT_PAGE_LOG2
92 } else {
93 page_log2
94 };
95 if !(8..=30).contains(&page_log2) {
96 return Err(format!("p64: page_log2 {page_log2} out of range"));
97 }
98 let page = 1usize << page_log2;
99 let plan = parse_safetensor_header(src)?;
100 if plan.tensors.is_empty() {
101 return Err("p64: Safetensors source contains no tensors".to_string());
102 }
103
104 let mut source_types = Vec::with_capacity(plan.tensors.len());
105 let mut roles = Vec::with_capacity(plan.tensors.len());
106 let mut n_layer = 0u32;
107 for tensor in &plan.tensors {
108 let ggml_type = safetensor_dtype_to_ggml(&tensor.dtype).ok_or_else(|| {
109 format!(
110 "p64: tensor '{}' dtype {} is not a supported high-fidelity source",
111 tensor.name, tensor.dtype
112 )
113 })?;
114 if !is_high_fidelity_ggml(ggml_type) {
115 return Err(format!(
116 "p64: tensor '{}' is low precision and cannot be transcoded",
117 tensor.name
118 ));
119 }
120 let element_count = tensor
121 .shape
122 .iter()
123 .try_fold(1usize, |count, dimension| count.checked_mul(*dimension))
124 .ok_or_else(|| format!("p64: tensor '{}' shape overflow", tensor.name))?;
125 let expected_bytes = element_count
126 .checked_mul(ggml_elem_bytes(ggml_type).ok_or("p64: unsupported element width")?)
127 .ok_or_else(|| format!("p64: tensor '{}' byte-size overflow", tensor.name))?;
128 if expected_bytes != tensor.byte_len() {
129 return Err(format!(
130 "p64: tensor '{}' declares {} bytes but shape/dtype requires {}",
131 tensor.name,
132 tensor.byte_len(),
133 expected_bytes
134 ));
135 }
136 let role = crate::tensor_roles::name_to_role(&tensor.name);
137 if let Some(mapped) = role {
138 if mapped.layer != P64_LAYER_GLOBAL {
139 n_layer = n_layer.max(mapped.layer as u32 + 1);
140 }
141 }
142 source_types.push(ggml_type);
143 roles.push(role);
144 }
145
146 let ternary_for = |position: usize| match policy {
147 SafetensorQuantization::Verbatim => false,
148 SafetensorQuantization::AllTernary => true,
149 SafetensorQuantization::FfnTernary => roles[position]
150 .map(|role| crate::tensor_roles::ternary_eligible(role.role))
151 .unwrap_or(false),
152 };
153
154 let mut string_table = vec![0u8];
155 let mut name_offsets = Vec::with_capacity(plan.tensors.len());
156 for tensor in &plan.tensors {
157 name_offsets.push(
158 u32::try_from(string_table.len()).map_err(|_| "p64: string table exceeds 4 GiB")?,
159 );
160 string_table.extend_from_slice(tensor.name.as_bytes());
161 string_table.push(0);
162 }
163
164 let hparams_offset = P64_WEIGHT_HEADER_BYTES;
165 let tensor_table_offset = align_up(hparams_offset + 64, 64);
166 let tensor_table_bytes = plan
167 .tensors
168 .len()
169 .checked_mul(P64_TENSOR_ENTRY_BYTES)
170 .ok_or("p64: tensor table overflow")?;
171 let string_table_offset = tensor_table_offset + tensor_table_bytes;
172 let manifold_table_offset = align_up(string_table_offset + string_table.len(), 64);
173 let manifold_count = n_layer as usize + 1;
174 let manifold_bytes = manifold_count
175 .checked_mul(P64_MANIFOLD_ENTRY_BYTES)
176 .ok_or("p64: manifold table overflow")?;
177 let tokenizer_offset = manifold_table_offset + manifold_bytes;
178 let checksum_offset = align_up(tokenizer_offset, 64);
179 let checksum_bytes = (plan.tensors.len() + 1)
180 .checked_mul(4)
181 .ok_or("p64: checksum table overflow")?;
182 let blob_region_offset = align_up(checksum_offset + checksum_bytes, page);
183
184 let mut entries = Vec::with_capacity(plan.tensors.len());
185 let mut cursor = blob_region_offset;
186 let mut largest_tensor_bytes = 0usize;
187 let mut total_tensor_bytes = 0usize;
188 let mut peak_working_bytes = 0usize;
189 for (position, tensor) in plan.tensors.iter().enumerate() {
190 let element_count = tensor.shape.iter().copied().product::<usize>();
191 let blob_size = if ternary_for(position) {
192 ternary_blob_len(element_count)
193 } else {
194 tensor.byte_len()
195 };
196 cursor = align_up(cursor, page);
197 let mapped = roles[position];
198 let role_id = mapped.map(|role| role.role).unwrap_or(P64_ROLE_UNKNOWN);
199 let manifold_idx = mapped
200 .filter(|role| role.layer != P64_LAYER_GLOBAL)
201 .map(|role| role.layer as u32)
202 .unwrap_or(n_layer);
203 let mut dimensions = [0u32; 4];
204 for (target, source) in dimensions.iter_mut().zip(tensor.shape.iter().take(4)) {
205 *target = u32::try_from(*source).map_err(|_| "p64: tensor dimension exceeds u32")?;
206 }
207 entries.push(P64TensorEntry {
208 name_offset: name_offsets[position],
209 role_id,
210 dtype: if ternary_for(position) {
211 GGML_TYPE_TERNARY_158 as u16
212 } else {
213 source_types[position] as u16
214 },
215 manifold_idx,
216 rank: tensor.shape.len().clamp(1, 4) as u32,
217 dimensions,
218 blob_offset: u32::try_from(cursor).map_err(|_| "p64: container exceeds 4 GiB")?,
219 blob_size: u32::try_from(blob_size).map_err(|_| "p64: tensor exceeds 4 GiB")?,
220 source_offset: tensor.begin as u64,
221 source_name_hash: crate::q_hash(&tensor.name),
222 alt_dtype: 0,
223 precision_views_mask: 0,
224 alt_blob_offset: 0,
225 });
226 cursor = cursor
227 .checked_add(blob_size)
228 .ok_or("p64: container size overflow")?;
229 largest_tensor_bytes = largest_tensor_bytes.max(blob_size);
230 total_tensor_bytes = total_tensor_bytes
231 .checked_add(blob_size)
232 .ok_or("p64: tensor byte total overflow")?;
233 peak_working_bytes = peak_working_bytes.max(tensor.byte_len());
234 }
235 if cursor > u32::MAX as usize {
236 return Err("p64: 32-bit relative-offset container exceeds 4 GiB".to_string());
237 }
238
239 let mut flags = P64_FLAG_LITTLE_ENDIAN | FORMAT_FLAG_RAW_TRANSCODE;
240 if !matches!(policy, SafetensorQuantization::Verbatim) {
241 flags |= FORMAT_FLAG_TERNARY;
242 }
243 let header = P64WeightHeader {
244 magic: P64_MAGIC,
245 version: P64_VERSION,
246 flags,
247 role_table_offset: 0,
248 tensor_table_offset: tensor_table_offset as u32,
249 tokenizer_offset: tokenizer_offset as u32,
250 hparams_offset: hparams_offset as u32,
251 string_table_offset: string_table_offset as u32,
252 checksum_offset: checksum_offset as u32,
253 manifold_table_offset: manifold_table_offset as u32,
254 tensor_count: entries.len() as u32,
255 page_size: page as u32,
256 reserved: [0; 20],
257 };
258 let hparams = P64HParams {
259 n_layer,
260 n_embd: 0,
261 n_head: 0,
262 n_kv_head: 0,
263 vocab_size: 0,
264 rope_freq_base: 0.0,
265 rope_scale: 0.0,
266 head_dim: 0,
267 head_dim_swa: 0,
268 sliding_window: 0,
269 shared_kv_layers: 0,
270 logit_softcap: 0.0,
271 architecture: 0,
272 arch_flags: 0,
273 reserved: [0; 8],
274 };
275
276 let mut metadata = vec![0u8; checksum_offset + checksum_bytes];
277 header.write_le(&mut metadata[..P64_WEIGHT_HEADER_BYTES]);
278 hparams.write_le(&mut metadata[hparams_offset..hparams_offset + 64]);
279 for (position, entry) in entries.iter().enumerate() {
280 let start = tensor_table_offset + position * P64_TENSOR_ENTRY_BYTES;
281 write_tensor_entry(entry, &mut metadata[start..start + P64_TENSOR_ENTRY_BYTES]);
282 }
283 metadata[string_table_offset..string_table_offset + string_table.len()]
284 .copy_from_slice(&string_table);
285 for layer in 0..manifold_count {
286 let coordinate = crate::modalities::manifold::ManifoldCoordinate10D::from_sequential_layer(
287 layer.min(n_layer as usize) as u32,
288 n_layer.max(1),
289 );
290 let start = manifold_table_offset + layer * P64_MANIFOLD_ENTRY_BYTES;
291 write_manifold_coordinate(
292 &coordinate,
293 &mut metadata[start..start + P64_MANIFOLD_ENTRY_BYTES],
294 );
295 }
296
297 let mut float_scratch = Vec::new();
298 for (position, tensor) in plan.tensors.iter().enumerate() {
299 let source_start = plan.data_start + tensor.begin;
300 let source_end = plan.data_start + tensor.end;
301 let crc = if ternary_for(position) {
302 let count = tensor.shape.iter().copied().product::<usize>();
303 decode_safetensor_to_f32(
304 &src[source_start..source_end],
305 source_types[position],
306 count,
307 &mut float_scratch,
308 );
309 if float_scratch.len() != count {
310 return Err(format!(
311 "p64: tensor '{}' decode was incomplete",
312 tensor.name
313 ));
314 }
315 let blob = ternary_blob(&float_scratch);
316 peak_working_bytes = peak_working_bytes.max(blob.len());
317 crc32c(&blob)
318 } else {
319 crc32c(&src[source_start..source_end])
320 };
321 let start = checksum_offset + 4 + position * 4;
322 metadata[start..start + 4].copy_from_slice(&crc.to_le_bytes());
323 }
324 let metadata_crc = crc32c(&metadata[..checksum_offset]);
325 metadata[checksum_offset..checksum_offset + 4].copy_from_slice(&metadata_crc.to_le_bytes());
326
327 out.write_all(&metadata)
328 .map_err(|error| error.to_string())?;
329 let zeros = [0u8; 4096];
330 let mut bytes_written = metadata.len();
331 for (position, tensor) in plan.tensors.iter().enumerate() {
332 let target = entries[position].blob_offset as usize;
333 while bytes_written < target {
334 let count = (target - bytes_written).min(zeros.len());
335 out.write_all(&zeros[..count])
336 .map_err(|error| error.to_string())?;
337 bytes_written += count;
338 }
339 let source_start = plan.data_start + tensor.begin;
340 let source_end = plan.data_start + tensor.end;
341 if ternary_for(position) {
342 let count = tensor.shape.iter().copied().product::<usize>();
343 decode_safetensor_to_f32(
344 &src[source_start..source_end],
345 source_types[position],
346 count,
347 &mut float_scratch,
348 );
349 let blob = ternary_blob(&float_scratch);
350 out.write_all(&blob).map_err(|error| error.to_string())?;
351 bytes_written += blob.len();
352 } else {
353 out.write_all(&src[source_start..source_end])
354 .map_err(|error| error.to_string())?;
355 bytes_written += source_end - source_start;
356 }
357 }
358
359 Ok(TranscodeReport {
360 n_tensors: entries.len(),
361 bytes_written,
362 largest_tensor_bytes,
363 total_tensor_bytes,
364 peak_working_bytes,
365 })
366}
367
368pub fn transcode_safetensor_to_p64<W: std::io::Write>(
369 src: &[u8],
370 page_log2: u16,
371 out: &mut W,
372) -> Result<TranscodeReport, String> {
373 transcode_safetensor_with_policy(src, page_log2, out, SafetensorQuantization::Verbatim)
374}
375
376pub fn transcode_safetensor_to_p64_ternary<W: std::io::Write>(
385 src: &[u8],
386 page_log2: u16,
387 out: &mut W,
388) -> Result<TranscodeReport, String> {
389 transcode_safetensor_with_policy(src, page_log2, out, SafetensorQuantization::AllTernary)
390}
391
392pub fn transcode_safetensor_to_p64_policy<W: std::io::Write>(
401 src: &[u8],
402 page_log2: u16,
403 out: &mut W,
404) -> Result<TranscodeReport, String> {
405 transcode_safetensor_with_policy(src, page_log2, out, SafetensorQuantization::FfnTernary)
406}
407
408pub fn transcode_safetensor_to_p64_ffn_ternary<W: std::io::Write>(
409 src: &[u8],
410 page_log2: u16,
411 out: &mut W,
412) -> Result<TranscodeReport, String> {
413 transcode_safetensor_to_p64_policy(src, page_log2, out)
414}
415
416pub fn transcode_safetensor_to_q42_ffn_ternary<W: std::io::Write>(
417 src: &[u8],
418 page_log2: u16,
419 out: &mut W,
420) -> Result<TranscodeReport, String> {
421 transcode_safetensor_to_p64_policy(src, page_log2, out)
422}
423
424pub const P64_ROLE_VISION_CONV2D: u16 = 0x80;
425pub const P64_ROLE_VISION_BN: u16 = 0x81;
426pub const P64_ROLE_VISION_FC: u16 = 0x82;
427
428#[derive(Debug, Clone)]
430pub struct RawVisionTensor {
431 pub name: String,
432 pub shape: Vec<usize>,
433 pub data: Vec<f32>,
434 pub role_id: u16,
435}
436
437pub fn transcode_vision_tensors_to_p64<W: std::io::Write>(
439 tensors: &[RawVisionTensor],
440 page_log2: u16,
441 out: &mut W,
442) -> Result<TranscodeReport, String> {
443 if tensors.is_empty() {
444 return Err("p64: Vision tensor list is empty".to_string());
445 }
446 let page_log2 = if page_log2 == 0 {
447 P64_DEFAULT_PAGE_LOG2
448 } else {
449 page_log2
450 };
451 let page = 1usize << page_log2;
452
453 let mut entries = Vec::with_capacity(tensors.len());
454 let mut names_blob = Vec::new();
455 let mut total_bytes = 0usize;
456 let mut max_working = 0usize;
457
458 let mut name_offsets = Vec::with_capacity(tensors.len());
459 for tensor in tensors {
460 let name_off = names_blob.len() as u32;
461 name_offsets.push(name_off);
462 names_blob.extend_from_slice(tensor.name.as_bytes());
463 names_blob.push(0);
464 }
465
466 let header_bytes = P64_WEIGHT_HEADER_BYTES; let hparams_offset = 64u32;
468 let hparams_bytes = 64usize;
469 let tensor_table_offset = (header_bytes + hparams_bytes) as u32; let entry_bytes = tensors.len() * P64_TENSOR_ENTRY_BYTES;
471 let string_table_offset = (tensor_table_offset as usize + entry_bytes) as u32;
472 let names_bytes = names_blob.len();
473 let manifold_table_offset = align_up((string_table_offset as usize) + names_bytes, 64) as u32;
474 let manifold_bytes = 64usize; let tokenizer_offset = (manifold_table_offset as usize + manifold_bytes) as u32;
476 let checksum_offset = tokenizer_offset;
477 let checksum_count = tensors.len() + 1;
478 let checksum_bytes = checksum_count * 4;
479 let meta_region_bytes = (checksum_offset as usize) + checksum_bytes;
480 let blob_region_start = align_up(meta_region_bytes, page);
481
482 let mut current_blob_offset = blob_region_start;
483 let mut tensor_crcs = Vec::with_capacity(tensors.len());
484
485 for (i, tensor) in tensors.iter().enumerate() {
486 let name_off = name_offsets[i];
487 let byte_len = tensor.data.len() * std::mem::size_of::<f32>();
488 max_working = max_working.max(byte_len);
489
490 let mut dimensions = [1u32; 4];
491 for (i, &dim) in tensor.shape.iter().take(4).enumerate() {
492 dimensions[i] = dim as u32;
493 }
494
495 let blob_off = current_blob_offset;
496 entries.push(P64TensorEntry {
497 name_offset: name_off,
498 role_id: tensor.role_id,
499 dtype: crate::safetensor::GGML_F32 as u16,
500 manifold_idx: 0,
501 rank: tensor.shape.len() as u32,
502 dimensions,
503 blob_offset: blob_off as u32,
504 blob_size: byte_len as u32,
505 source_offset: 0,
506 source_name_hash: crate::q_hash(&tensor.name),
507 alt_dtype: 0,
508 precision_views_mask: 0,
509 alt_blob_offset: 0,
510 });
511
512 let mut raw_bytes = Vec::with_capacity(byte_len);
514 for &val in &tensor.data {
515 raw_bytes.extend_from_slice(&val.to_le_bytes());
516 }
517 let crc = crate::container_10d::crc32c::crc32c(&raw_bytes);
518 tensor_crcs.push(crc);
519
520 current_blob_offset = align_up(current_blob_offset + byte_len, page);
521 total_bytes += byte_len;
522 }
523
524 let mut meta_buf = vec![0u8; checksum_offset as usize];
525
526 let hdr = P64WeightHeader {
527 magic: P64_MAGIC,
528 version: P64_VERSION,
529 flags: P64_FLAG_LITTLE_ENDIAN,
530 role_table_offset: tensor_table_offset,
531 tensor_table_offset,
532 tokenizer_offset,
533 hparams_offset,
534 string_table_offset,
535 checksum_offset,
536 manifold_table_offset,
537 tensor_count: entries.len() as u32,
538 page_size: page as u32,
539 reserved: [0u8; 20],
540 };
541 hdr.write_le(&mut meta_buf[0..64]);
542
543 let hparams = P64HParams {
545 n_layer: 0,
546 n_embd: 0,
547 n_head: 0,
548 n_kv_head: 0,
549 vocab_size: 0,
550 rope_freq_base: 10000.0,
551 rope_scale: 1.0,
552 head_dim: 0,
553 head_dim_swa: 0,
554 sliding_window: 0,
555 shared_kv_layers: 0,
556 logit_softcap: 0.0,
557 architecture: 0,
558 arch_flags: 0,
559 reserved: [0u8; 8],
560 };
561 hparams.write_le(&mut meta_buf[64..128]);
562
563 for (i, entry) in entries.iter().enumerate() {
565 let off = (tensor_table_offset as usize) + i * P64_TENSOR_ENTRY_BYTES;
566 write_tensor_entry(entry, &mut meta_buf[off..off + P64_TENSOR_ENTRY_BYTES]);
567 }
568
569 meta_buf[string_table_offset as usize..(string_table_offset as usize) + names_bytes]
571 .copy_from_slice(&names_blob);
572
573 let meta_crc = crate::container_10d::crc32c::crc32c(&meta_buf);
575
576 out.write_all(&meta_buf).map_err(|e| e.to_string())?;
577
578 let mut checksum_buf = vec![0u8; checksum_bytes];
580 checksum_buf[0..4].copy_from_slice(&meta_crc.to_le_bytes());
581 for (i, &crc) in tensor_crcs.iter().enumerate() {
582 let off = 4 + i * 4;
583 checksum_buf[off..off + 4].copy_from_slice(&crc.to_le_bytes());
584 }
585 out.write_all(&checksum_buf).map_err(|e| e.to_string())?;
586
587 let pad_to_blobs = blob_region_start - meta_region_bytes;
589 if pad_to_blobs > 0 {
590 out.write_all(&vec![0u8; pad_to_blobs])
591 .map_err(|e| e.to_string())?;
592 }
593
594 let mut written_bytes = blob_region_start;
595 for tensor in tensors {
596 let mut raw_bytes = Vec::with_capacity(tensor.data.len() * 4);
597 for &val in &tensor.data {
598 raw_bytes.extend_from_slice(&val.to_le_bytes());
599 }
600 out.write_all(&raw_bytes).map_err(|e| e.to_string())?;
601 written_bytes += raw_bytes.len();
602
603 let pad = align_up(raw_bytes.len(), page) - raw_bytes.len();
604 if pad > 0 {
605 out.write_all(&vec![0u8; pad]).map_err(|e| e.to_string())?;
606 written_bytes += pad;
607 }
608 }
609
610 Ok(TranscodeReport {
611 n_tensors: tensors.len(),
612 bytes_written: written_bytes,
613 largest_tensor_bytes: max_working,
614 total_tensor_bytes: total_bytes,
615 peak_working_bytes: max_working,
616 })
617}