528 lines
17 KiB
Rust
528 lines
17 KiB
Rust
use memmap2::{Mmap, MmapOptions};
|
|
use std::collections::HashMap;
|
|
use std::fs::File;
|
|
use std::path::{Path, PathBuf};
|
|
|
|
const MAGIC: u32 = 0x4655_4747;
|
|
const MAX_COUNT: u64 = 1_000_000;
|
|
const MAX_DIMS: usize = 8;
|
|
const MAX_DEPTH: u8 = 8;
|
|
|
|
pub(super) const F32: u32 = 0;
|
|
pub(super) const F16: u32 = 1;
|
|
pub(super) const Q4_0: u32 = 2;
|
|
pub(super) const Q8_0: u32 = 8;
|
|
pub(super) const Q2_K: u32 = 10;
|
|
pub(super) const Q4_K: u32 = 12;
|
|
pub(super) const Q5_K: u32 = 13;
|
|
pub(super) const Q6_K: u32 = 14;
|
|
pub(super) const IQ2_XXS: u32 = 16;
|
|
pub(super) const I32: u32 = 26;
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub(super) enum Value {
|
|
U32(u32),
|
|
I32(i32),
|
|
U64(u64),
|
|
I64(i64),
|
|
F32(f32),
|
|
F64(f64),
|
|
Bool(bool),
|
|
Bytes(Vec<u8>),
|
|
U32s(Vec<u32>),
|
|
F32s(Vec<f32>),
|
|
Strings(Vec<Vec<u8>>),
|
|
Other,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub(super) struct Tensor {
|
|
pub(super) kind: u32,
|
|
pub(super) dims: Vec<u64>,
|
|
pub(super) offset: u64,
|
|
pub(super) bytes: u64,
|
|
}
|
|
|
|
pub(super) struct Gguf {
|
|
path: PathBuf,
|
|
map: Mmap,
|
|
data_offset: u64,
|
|
max_tensor_bytes: u64,
|
|
pub(super) metadata: HashMap<String, Value>,
|
|
pub(super) tensors: HashMap<String, Tensor>,
|
|
}
|
|
|
|
impl Gguf {
|
|
pub(super) fn open(path: &Path) -> Result<Self, String> {
|
|
let file =
|
|
File::open(path).map_err(|error| format!("Cannot open {}: {error}", path.display()))?;
|
|
if file.metadata().map_err(|error| error.to_string())?.len() < 32 {
|
|
return Err(format!("{} is too small to be GGUF", path.display()));
|
|
}
|
|
// SAFETY: managed model artifacts are opened read-only and are never mutated
|
|
// while a Model owns this mapping.
|
|
let map = unsafe { MmapOptions::new().map(&file) }
|
|
.map_err(|error| format!("Cannot map {}: {error}", path.display()))?;
|
|
let mut cursor = Cursor::new(&map);
|
|
if cursor.u32()? != MAGIC {
|
|
return Err(format!("{} is not a GGUF file", path.display()));
|
|
}
|
|
if cursor.u32()? != 3 {
|
|
return Err(format!("{} is not GGUF v3", path.display()));
|
|
}
|
|
let tensor_count = cursor.u64()?;
|
|
let metadata_count = cursor.u64()?;
|
|
if tensor_count > MAX_COUNT || metadata_count > MAX_COUNT {
|
|
return Err("GGUF directory count is too large".into());
|
|
}
|
|
|
|
let mut metadata = HashMap::with_capacity(metadata_count as usize);
|
|
let mut alignment = 32_u64;
|
|
for _ in 0..metadata_count {
|
|
let key = String::from_utf8(cursor.string()?.to_vec())
|
|
.map_err(|_| "GGUF metadata key is not UTF-8".to_owned())?;
|
|
let kind = cursor.u32()?;
|
|
let value = cursor.value(kind, &key, 0)?;
|
|
if key == "general.alignment"
|
|
&& let Value::U32(value) = value
|
|
{
|
|
alignment = u64::from(value);
|
|
}
|
|
if metadata.insert(key.clone(), value).is_some() {
|
|
return Err(format!("duplicate GGUF metadata key: {key}"));
|
|
}
|
|
}
|
|
if alignment == 0 {
|
|
return Err("GGUF alignment is zero".into());
|
|
}
|
|
|
|
let mut tensors = HashMap::with_capacity(tensor_count as usize);
|
|
for _ in 0..tensor_count {
|
|
let name = String::from_utf8(cursor.string()?.to_vec())
|
|
.map_err(|_| "GGUF tensor name is not UTF-8".to_owned())?;
|
|
let ndim = cursor.u32()? as usize;
|
|
if ndim == 0 || ndim > MAX_DIMS {
|
|
return Err(format!(
|
|
"tensor {name} has unsupported dimension count {ndim}"
|
|
));
|
|
}
|
|
let mut dims = Vec::with_capacity(ndim);
|
|
let mut elements = 1_u64;
|
|
for _ in 0..ndim {
|
|
let dim = cursor.u64()?;
|
|
elements = elements
|
|
.checked_mul(dim)
|
|
.ok_or_else(|| format!("tensor {name} element count overflows"))?;
|
|
dims.push(dim);
|
|
}
|
|
let kind = cursor.u32()?;
|
|
let relative_offset = cursor.u64()?;
|
|
let (block_elements, block_bytes) = tensor_type(kind)
|
|
.ok_or_else(|| format!("tensor {name} uses unsupported GGUF type {kind}"))?;
|
|
let blocks = elements
|
|
.checked_add(block_elements - 1)
|
|
.ok_or_else(|| format!("tensor {name} size overflows"))?
|
|
/ block_elements;
|
|
let bytes = blocks
|
|
.checked_mul(block_bytes)
|
|
.ok_or_else(|| format!("tensor {name} size overflows"))?;
|
|
if tensors
|
|
.insert(
|
|
name.clone(),
|
|
Tensor {
|
|
kind,
|
|
dims,
|
|
offset: relative_offset,
|
|
bytes,
|
|
},
|
|
)
|
|
.is_some()
|
|
{
|
|
return Err(format!("duplicate GGUF tensor: {name}"));
|
|
}
|
|
}
|
|
|
|
let data_start = align(cursor.position() as u64, alignment)?;
|
|
let mut max_tensor_bytes = 0;
|
|
for (name, tensor) in &mut tensors {
|
|
tensor.offset = data_start
|
|
.checked_add(tensor.offset)
|
|
.ok_or_else(|| format!("tensor {name} offset overflows"))?;
|
|
let end = tensor
|
|
.offset
|
|
.checked_add(tensor.bytes)
|
|
.ok_or_else(|| format!("tensor {name} end overflows"))?;
|
|
if end > map.len() as u64 {
|
|
return Err(format!("tensor {name} points outside the GGUF file"));
|
|
}
|
|
max_tensor_bytes = max_tensor_bytes.max(tensor.bytes);
|
|
}
|
|
|
|
Ok(Self {
|
|
path: path.to_owned(),
|
|
map,
|
|
data_offset: data_start,
|
|
max_tensor_bytes,
|
|
metadata,
|
|
tensors,
|
|
})
|
|
}
|
|
|
|
pub(super) fn path(&self) -> &Path {
|
|
&self.path
|
|
}
|
|
|
|
pub(super) fn len(&self) -> u64 {
|
|
self.map.len() as u64
|
|
}
|
|
|
|
pub(super) fn map_ptr(&self) -> *const u8 {
|
|
self.map.as_ptr()
|
|
}
|
|
|
|
pub(super) fn data_offset(&self) -> u64 {
|
|
self.data_offset
|
|
}
|
|
|
|
pub(super) fn max_tensor_bytes(&self) -> u64 {
|
|
self.max_tensor_bytes
|
|
}
|
|
|
|
pub(super) fn tensor(&self, name: &str) -> Result<&Tensor, String> {
|
|
self.tensors
|
|
.get(name)
|
|
.ok_or_else(|| format!("required tensor is missing: {name}"))
|
|
}
|
|
|
|
pub(super) fn tensor_data(&self, name: &str) -> Result<&[u8], String> {
|
|
let tensor = self.tensor(name)?;
|
|
let start = tensor.offset as usize;
|
|
let end = start + tensor.bytes as usize;
|
|
Ok(&self.map[start..end])
|
|
}
|
|
|
|
pub(super) fn u32(&self, key: &str) -> Result<u32, String> {
|
|
match self.metadata.get(key) {
|
|
Some(Value::U32(value)) => Ok(*value),
|
|
_ => Err(format!("required uint32 metadata key is missing: {key}")),
|
|
}
|
|
}
|
|
|
|
pub(super) fn u64(&self, key: &str) -> Result<u64, String> {
|
|
match self.metadata.get(key) {
|
|
Some(Value::U64(value)) => Ok(*value),
|
|
Some(Value::U32(value)) => Ok(u64::from(*value)),
|
|
_ => Err(format!("required integer metadata key is missing: {key}")),
|
|
}
|
|
}
|
|
|
|
pub(super) fn f32(&self, key: &str) -> Result<f32, String> {
|
|
match self.metadata.get(key) {
|
|
Some(Value::F32(value)) => Ok(*value),
|
|
Some(Value::F64(value)) => Ok(*value as f32),
|
|
Some(Value::U32(value)) => Ok(*value as f32),
|
|
Some(Value::I32(value)) => Ok(*value as f32),
|
|
_ => Err(format!("required numeric metadata key is missing: {key}")),
|
|
}
|
|
}
|
|
|
|
pub(super) fn token_id(&self, key: &str) -> Result<i32, String> {
|
|
let value = match self.metadata.get(key) {
|
|
Some(Value::U32(value)) => i64::from(*value),
|
|
Some(Value::I32(value)) => i64::from(*value),
|
|
Some(Value::U64(value)) => i64::try_from(*value).unwrap_or(-1),
|
|
Some(Value::I64(value)) => *value,
|
|
_ => -1,
|
|
};
|
|
i32::try_from(value)
|
|
.ok()
|
|
.filter(|value| *value >= 0)
|
|
.ok_or_else(|| format!("required tokenizer token ID is missing: {key}"))
|
|
}
|
|
|
|
pub(super) fn boolean(&self, key: &str) -> Result<bool, String> {
|
|
match self.metadata.get(key) {
|
|
Some(Value::Bool(value)) => Ok(*value),
|
|
_ => Err(format!("required boolean metadata key is missing: {key}")),
|
|
}
|
|
}
|
|
|
|
pub(super) fn bytes(&self, key: &str) -> Result<&[u8], String> {
|
|
match self.metadata.get(key) {
|
|
Some(Value::Bytes(value)) => Ok(value),
|
|
_ => Err(format!("required string metadata key is missing: {key}")),
|
|
}
|
|
}
|
|
|
|
pub(super) fn u32s(&self, key: &str) -> Result<&[u32], String> {
|
|
match self.metadata.get(key) {
|
|
Some(Value::U32s(value)) => Ok(value),
|
|
_ => Err(format!(
|
|
"required integer-array metadata key is missing: {key}"
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(super) fn f32s(&self, key: &str) -> Result<&[f32], String> {
|
|
match self.metadata.get(key) {
|
|
Some(Value::F32s(value)) => Ok(value),
|
|
_ => Err(format!(
|
|
"required float-array metadata key is missing: {key}"
|
|
)),
|
|
}
|
|
}
|
|
|
|
pub(super) fn strings(&self, key: &str) -> Result<&[Vec<u8>], String> {
|
|
match self.metadata.get(key) {
|
|
Some(Value::Strings(value)) => Ok(value),
|
|
_ => Err(format!(
|
|
"required string-array metadata key is missing: {key}"
|
|
)),
|
|
}
|
|
}
|
|
}
|
|
|
|
fn tensor_type(kind: u32) -> Option<(u64, u64)> {
|
|
Some(match kind {
|
|
0 => (1, 4),
|
|
1 => (1, 2),
|
|
2 => (32, 18),
|
|
3 => (32, 20),
|
|
6 => (32, 22),
|
|
7 => (32, 24),
|
|
8 => (32, 34),
|
|
9 => (32, 40),
|
|
10 => (256, 84),
|
|
11 => (256, 110),
|
|
12 => (256, 144),
|
|
13 => (256, 176),
|
|
14 => (256, 210),
|
|
15 => (256, 292),
|
|
16 => (256, 66),
|
|
17 => (256, 74),
|
|
18 => (256, 98),
|
|
19 => (256, 110),
|
|
20 => (256, 50),
|
|
21 => (256, 110),
|
|
22 => (256, 82),
|
|
23 => (256, 136),
|
|
24 => (1, 1),
|
|
25 => (1, 2),
|
|
26 => (1, 4),
|
|
27 => (1, 8),
|
|
28 => (1, 8),
|
|
29 => (256, 56),
|
|
30 => (1, 2),
|
|
_ => return None,
|
|
})
|
|
}
|
|
|
|
fn align(value: u64, alignment: u64) -> Result<u64, String> {
|
|
let remainder = value % alignment;
|
|
if remainder == 0 {
|
|
Ok(value)
|
|
} else {
|
|
value
|
|
.checked_add(alignment - remainder)
|
|
.ok_or_else(|| "GGUF alignment overflows".into())
|
|
}
|
|
}
|
|
|
|
struct Cursor<'a> {
|
|
bytes: &'a [u8],
|
|
position: usize,
|
|
}
|
|
|
|
impl<'a> Cursor<'a> {
|
|
fn new(bytes: &'a [u8]) -> Self {
|
|
Self { bytes, position: 0 }
|
|
}
|
|
|
|
fn position(&self) -> usize {
|
|
self.position
|
|
}
|
|
|
|
fn take(&mut self, count: u64) -> Result<&'a [u8], String> {
|
|
let count = usize::try_from(count).map_err(|_| "GGUF value is too large")?;
|
|
let end = self
|
|
.position
|
|
.checked_add(count)
|
|
.filter(|end| *end <= self.bytes.len())
|
|
.ok_or_else(|| format!("truncated GGUF at byte {}", self.position))?;
|
|
let value = &self.bytes[self.position..end];
|
|
self.position = end;
|
|
Ok(value)
|
|
}
|
|
|
|
fn u8(&mut self) -> Result<u8, String> {
|
|
Ok(self.take(1)?[0])
|
|
}
|
|
|
|
fn u16(&mut self) -> Result<u16, String> {
|
|
Ok(u16::from_le_bytes(self.take(2)?.try_into().unwrap()))
|
|
}
|
|
|
|
fn u32(&mut self) -> Result<u32, String> {
|
|
Ok(u32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
|
}
|
|
|
|
fn i32(&mut self) -> Result<i32, String> {
|
|
Ok(i32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
|
}
|
|
|
|
fn u64(&mut self) -> Result<u64, String> {
|
|
Ok(u64::from_le_bytes(self.take(8)?.try_into().unwrap()))
|
|
}
|
|
|
|
fn i64(&mut self) -> Result<i64, String> {
|
|
Ok(i64::from_le_bytes(self.take(8)?.try_into().unwrap()))
|
|
}
|
|
|
|
fn f32(&mut self) -> Result<f32, String> {
|
|
Ok(f32::from_le_bytes(self.take(4)?.try_into().unwrap()))
|
|
}
|
|
|
|
fn f64(&mut self) -> Result<f64, String> {
|
|
Ok(f64::from_le_bytes(self.take(8)?.try_into().unwrap()))
|
|
}
|
|
|
|
fn string(&mut self) -> Result<&'a [u8], String> {
|
|
let len = self.u64()?;
|
|
self.take(len)
|
|
}
|
|
|
|
fn value(&mut self, kind: u32, key: &str, depth: u8) -> Result<Value, String> {
|
|
if depth > MAX_DEPTH {
|
|
return Err("GGUF metadata arrays are nested too deeply".into());
|
|
}
|
|
Ok(match kind {
|
|
0 => {
|
|
self.u8()?;
|
|
Value::Other
|
|
}
|
|
1 => {
|
|
self.u8()?;
|
|
Value::Other
|
|
}
|
|
2 | 3 => {
|
|
self.u16()?;
|
|
Value::Other
|
|
}
|
|
4 => Value::U32(self.u32()?),
|
|
5 => Value::I32(self.i32()?),
|
|
6 => Value::F32(self.f32()?),
|
|
7 => Value::Bool(self.u8()? != 0),
|
|
8 => Value::Bytes(self.string()?.to_vec()),
|
|
9 => self.array(key, depth + 1)?,
|
|
10 => Value::U64(self.u64()?),
|
|
11 => Value::I64(self.i64()?),
|
|
12 => Value::F64(self.f64()?),
|
|
_ => return Err(format!("unknown GGUF metadata type {kind}")),
|
|
})
|
|
}
|
|
|
|
fn array(&mut self, key: &str, depth: u8) -> Result<Value, String> {
|
|
let item = self.u32()?;
|
|
let len = self.u64()?;
|
|
if len > MAX_COUNT {
|
|
return Err(format!("GGUF metadata array is too large: {key}"));
|
|
}
|
|
let keep = matches!(
|
|
key,
|
|
"tokenizer.ggml.tokens"
|
|
| "tokenizer.ggml.merges"
|
|
| "deepseek4.attention.compress_ratios"
|
|
| "deepseek4.swiglu_clamp_exp"
|
|
| "deepseek4.dspark.target_layer_ids"
|
|
| "deepseek4.dspark_target_layer_ids"
|
|
| "dspark.target_layer_ids"
|
|
);
|
|
if !keep {
|
|
for _ in 0..len {
|
|
self.value(item, key, depth)?;
|
|
}
|
|
return Ok(Value::Other);
|
|
}
|
|
match item {
|
|
4 | 5 => {
|
|
let mut values = Vec::with_capacity(len as usize);
|
|
for _ in 0..len {
|
|
let value = if item == 4 {
|
|
self.u32()?
|
|
} else {
|
|
self.i32()?
|
|
.try_into()
|
|
.map_err(|_| format!("negative value in {key}"))?
|
|
};
|
|
values.push(value);
|
|
}
|
|
Ok(Value::U32s(values))
|
|
}
|
|
6 | 12 => {
|
|
let mut values = Vec::with_capacity(len as usize);
|
|
for _ in 0..len {
|
|
values.push(if item == 6 {
|
|
self.f32()?
|
|
} else {
|
|
self.f64()? as f32
|
|
});
|
|
}
|
|
Ok(Value::F32s(values))
|
|
}
|
|
8 => {
|
|
let mut values = Vec::with_capacity(len as usize);
|
|
for _ in 0..len {
|
|
values.push(self.string()?.to_vec());
|
|
}
|
|
Ok(Value::Strings(values))
|
|
}
|
|
_ => Err(format!("unsupported array type {item} for {key}")),
|
|
}
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use std::fs;
|
|
use std::time::{SystemTime, UNIX_EPOCH};
|
|
|
|
#[test]
|
|
fn maps_metadata_and_tensor_payload_without_copying_weights() {
|
|
let path = std::env::temp_dir().join(format!(
|
|
"ds4-server-gguf-{}",
|
|
SystemTime::now()
|
|
.duration_since(UNIX_EPOCH)
|
|
.unwrap()
|
|
.as_nanos()
|
|
));
|
|
let mut bytes = Vec::new();
|
|
bytes.extend(MAGIC.to_le_bytes());
|
|
bytes.extend(3_u32.to_le_bytes());
|
|
bytes.extend(1_u64.to_le_bytes());
|
|
bytes.extend(1_u64.to_le_bytes());
|
|
push_string(&mut bytes, b"general.architecture");
|
|
bytes.extend(8_u32.to_le_bytes());
|
|
push_string(&mut bytes, b"deepseek4");
|
|
push_string(&mut bytes, b"weight");
|
|
bytes.extend(1_u32.to_le_bytes());
|
|
bytes.extend(1_u64.to_le_bytes());
|
|
bytes.extend(F32.to_le_bytes());
|
|
bytes.extend(0_u64.to_le_bytes());
|
|
bytes.resize(bytes.len().div_ceil(32) * 32, 0);
|
|
bytes.extend(1_f32.to_le_bytes());
|
|
fs::write(&path, bytes).unwrap();
|
|
|
|
let model = Gguf::open(&path).unwrap();
|
|
assert_eq!(model.bytes("general.architecture").unwrap(), b"deepseek4");
|
|
assert_eq!(model.tensor("weight").unwrap().dims, [1]);
|
|
assert_eq!(model.tensor_data("weight").unwrap(), 1_f32.to_le_bytes());
|
|
fs::remove_file(path).unwrap();
|
|
}
|
|
|
|
fn push_string(bytes: &mut Vec<u8>, value: &[u8]) {
|
|
bytes.extend((value.len() as u64).to_le_bytes());
|
|
bytes.extend(value);
|
|
}
|
|
}
|