diff --git a/CHANGELOG.md b/CHANGELOG.md index 2fcd89c..955c6a5 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -24,13 +24,17 @@ All notable changes to this project are documented in this file. returned. Unreferenced on-disk shards are reported via `index:unreferenced_shards` rather than rejected. Directory listings and tensor records are sorted so repeated inspection is byte-identical. -- **Optional `mmap` feature (#45 option 1):** `load_gguf_mmap` maps a GGUF with - `memmap2` 0.9.11 instead of `fs::read` into a `Vec`. Default builds stay - zero-dep (`default = []`). Packed CPU dequant for **Q8_0**, **Q5_K**, - **Q6_K**, and the internal **IQ3_M block** layout +- **Optional `mmap` feature (#45 option 1 / RM-367):** `load_gguf_mmap` maps a + GGUF with `memmap2` 0.9.11 instead of `fs::read` into a `Vec`. Default + builds stay zero-dep (`default = []`). Packed CPU dequant for **Q8_0**, + **Q5_K**, **Q6_K**, and the internal **IQ3_M block** layout (`GGML_TYPE_IQ3_M_BLOCK = 0x4949334D`, 111 bytes / 256 values). Wire type **31** remains historical **Q4_0_4_4** (`DType::Other(31)`). CUDA - host-register is still out of scope. + host-register is still out of scope. `GgufLayoutMmap::directory_matches` + compares the full tensor directory (not just count/architecture). CI covers + mmap↔owned packed dequant, a tensor past the first OS page, and a **sparse + 2 GiB** mapping that never calls `fs::read`. Real on-disk multi-GB pilots + remain `#[ignore]` behind `ENGRAM_GGUF`. ### Changed diff --git a/README.md b/README.md index f5ebd6f..c90cb92 100644 --- a/README.md +++ b/README.md @@ -272,7 +272,7 @@ cargo test --all-features cargo llvm-cov --all-targets --all-features --locked --lcov --output-path lcov.info ``` -`load_gguf` reads and retains the complete file. For multi-GB checkpoints use `cargo test --features mmap` / `load_gguf_mmap`. Real GGUF pilots require local files (`ENGRAM_GGUF`) and are `#[ignore]` so CI stays green without them. CUDA host-register is still out of scope. +`load_gguf` reads and retains the complete file. For multi-GB checkpoints use `cargo test --features mmap` / `load_gguf_mmap`. CI mmap tests include packed Q8_0/Q5_K/Q6_K/IQ3_M dequant from the mapping and a sparse 2 GiB file that is mapped without `fs::read`. Real on-disk GGUF pilots still require local files (`ENGRAM_GGUF`) and are `#[ignore]`. CUDA host-register is still out of scope. ```bash ENGRAM_GGUF=~/.models/gguf/.../model.gguf ENGRAM_EXPECT_MOE=1 \ diff --git a/src/gguf/dequant.rs b/src/gguf/dequant.rs index 914099a..92d16fc 100644 --- a/src/gguf/dequant.rs +++ b/src/gguf/dequant.rs @@ -146,6 +146,7 @@ fn fail_closed_dequant(dtype: DType) -> ParserError { pub(crate) fn dequantize_row_q8_0(row: &[u8], width: usize) -> Result> { require_multiple(width, Q8_0_BLOCK, "Q8_0")?; + expect_row_len(row, width, Q8_0_BLOCK, Q8_0_BYTES, "Q8_0")?; let mut out = Vec::with_capacity(width); let (blocks, _) = row.as_chunks::(); for block in blocks { @@ -159,6 +160,7 @@ pub(crate) fn dequantize_row_q8_0(row: &[u8], width: usize) -> Result> pub(crate) fn dequantize_row_q5_k(row: &[u8], width: usize) -> Result> { require_multiple(width, K_BLOCK, "Q5_K")?; + expect_row_len(row, width, K_BLOCK, Q5_K_BYTES, "Q5_K")?; let mut out = Vec::with_capacity(width); let (blocks, _) = row.as_chunks::(); for block in blocks { @@ -537,12 +539,17 @@ mod tests { fn dequant_error_paths() { assert!(dequantize_q8_0(&[0u8; 34], &[31]).is_err()); assert!(dequantize_q8_0(&[0u8; 33], &[32]).is_err()); + assert!(dequantize_row_q8_0(&[0u8; 33], 32).is_err()); + assert!(dequantize_row_q5_k(&[0u8; 175], 256).is_err()); assert!(dequantize_iq3_m(&[0u8; 111], &[255]).is_err()); assert!(dequantize_iq3_m(&[0u8; 110], &[256]).is_err()); assert!(dequantize_q5_k(&[0u8; 176], &[]).is_err()); assert!(dequantize_q8_0(&[], &[0]).is_err()); assert!(dequantize_packed(DType::F32, &[0u8; 4], &[1]).is_err()); assert!(dequantize_packed(DType::Other(31), &[0u8; 18], &[32]).is_err()); + assert!(dequantize_packed(DType::Q5_K, &[0u8; 176], &[256]).is_ok()); + assert!(dequantize_packed(DType::Q6_K, &[0u8; 210], &[256]).is_ok()); + assert!(dequantize_packed(DType::IQ3_M_BLOCK, &[0u8; 111], &[256]).is_ok()); } #[test] diff --git a/src/gguf/map.rs b/src/gguf/map.rs index e72e830..cf02ffa 100644 --- a/src/gguf/map.rs +++ b/src/gguf/map.rs @@ -183,11 +183,27 @@ impl GgufLayoutMmap { } /// Compare directory metadata with an owned [`GgufLayout`] of the same file. + /// + /// Checks alignment, tensor-data offset, architecture, and every tensor's + /// name, shape, dtype, packed length, and offsets. Does not compare + /// payload bytes (callers that need that should use + /// [`Self::tensor_bytes`] against [`GgufLayout::tensor_bytes`]). pub fn directory_matches(&self, owned: &GgufLayout) -> bool { self.alignment == owned.alignment && self.tensor_data_offset == owned.tensor_data_offset && self.tensors.len() == owned.tensors.len() && self.metadata.architecture() == owned.metadata.architecture() + && self.tensors.iter().all(|(name, mapped)| { + owned.tensors.get(name).is_some_and(|owned_t| { + mapped.dims == owned_t.dims + && mapped.dtype == owned_t.dtype + && mapped.ggml_type == owned_t.ggml_type + && mapped.n_elements == owned_t.n_elements + && mapped.byte_len == owned_t.byte_len + && mapped.relative_offset == owned_t.relative_offset + && mapped.absolute_offset == owned_t.absolute_offset + }) + }) } } diff --git a/tests/common/mod.rs b/tests/common/mod.rs index d24d0be..d8efce6 100644 --- a/tests/common/mod.rs +++ b/tests/common/mod.rs @@ -16,7 +16,11 @@ pub const VT_STRING: u32 = 8; pub const GGML_F32: u32 = 0; pub const GGML_Q8_0: u32 = 8; pub const GGML_Q4_K: u32 = 12; +pub const GGML_Q5_K: u32 = 13; +pub const GGML_Q6_K: u32 = 14; pub const GGML_IQ3_S: u32 = 21; +/// Internal 111-byte IQ3_M block id (not GGUF wire type 31). +pub const GGML_IQ3_M_BLOCK: u32 = 0x4949_334D; pub enum KvValue { U32(u32), diff --git a/tests/mmap_gguf.rs b/tests/mmap_gguf.rs index 94b91c8..a268a9a 100644 --- a/tests/mmap_gguf.rs +++ b/tests/mmap_gguf.rs @@ -1,18 +1,31 @@ // SPDX-License-Identifier: MIT OR Apache-2.0 -//! mmap-backed GGUF reader vs owned `load_gguf` / `parse_bytes` on a tiny file. +//! mmap-backed GGUF reader vs owned `load_gguf` / `parse_bytes`. +//! +//! Covers tiny-file parity, packed K-quant dequant from the mapping, +//! a tensor that starts past the first OS page, and a sparse multi-GiB +//! file that is mapped without `fs::read`. #![cfg(feature = "mmap")] mod common; use common::*; -use engram_parser::{load_gguf, load_gguf_mmap, parse_bytes}; +use engram_parser::{ + DType, dequantize_iq3_m, dequantize_packed, dequantize_q5_k, dequantize_q6_k, dequantize_q8_0, + load_gguf, load_gguf_mmap, os_page_size, parse_bytes, +}; use std::fs; -use std::path::PathBuf; +use std::fs::OpenOptions; +use std::os::unix::fs::MetadataExt; +use std::path::{Path, PathBuf}; use std::process; use std::time::{SystemTime, UNIX_EPOCH}; +/// 2 GiB sparse payload — large enough to be a multi-GB checkpoint stand-in +/// without allocating RSS (ext4 hole). +const SPARSE_PAYLOAD_BYTES: u64 = 2 * 1024 * 1024 * 1024; + fn write_temp_gguf(bytes: &[u8]) -> PathBuf { let nanos = SystemTime::now() .duration_since(UNIX_EPOCH) @@ -24,6 +37,83 @@ fn write_temp_gguf(bytes: &[u8]) -> PathBuf { path } +fn f16_one() -> [u8; 2] { + 0x3C00u16.to_le_bytes() +} + +fn q8_0_ones_block() -> Vec { + let mut block = vec![0u8; 34]; + block[0..2].copy_from_slice(&f16_one()); + for q in &mut block[2..] { + *q = 1; + } + block +} + +fn q5_k_ones_block() -> Vec { + let mut block = vec![0u8; 176]; + block[0..2].copy_from_slice(&f16_one()); + for b in &mut block[4..16] { + *b = 0x01; + } + for b in &mut block[48..] { + *b = 0x11; + } + block +} + +fn q6_k_minus_32_block() -> Vec { + let mut block = vec![0u8; 210]; + for b in &mut block[192..208] { + *b = 1; + } + block[208..210].copy_from_slice(&f16_one()); + block +} + +fn iq3_m_scaled_block() -> Vec { + let mut block = vec![0u8; 111]; + block[0..2].copy_from_slice(&f16_one()); + for b in &mut block[34..98] { + *b = 0x55; + } + block[98] = 0x01; + block +} + +/// Header-only GGUF whose F32 payload is a sparse hole of `n_elements * 4` bytes. +fn write_sparse_f32_gguf(path: &Path, n_elements: u64) -> u64 { + let mut out = Vec::new(); + out.extend_from_slice(&GGUF_MAGIC); + push_u32(&mut out, GGUF_VERSION); + push_u64(&mut out, 1); + push_u64(&mut out, 2); + push_kv_u32(&mut out, "general.alignment", ALIGNMENT); + push_kv_string(&mut out, "general.architecture", "olmoe"); + push_string(&mut out, "sparse.weight"); + push_u32(&mut out, 1); + push_u64(&mut out, n_elements); + push_u32(&mut out, GGML_F32); + push_u64(&mut out, 0); + while !out.len().is_multiple_of(ALIGNMENT as usize) { + out.push(0); + } + let header_len = out.len() as u64; + let payload = n_elements + .checked_mul(4) + .expect("sparse payload byte-length overflow"); + let file_len = header_len + .checked_add(payload) + .expect("sparse file length overflow"); + fs::write(path, &out).expect("write sparse header"); + let file = OpenOptions::new() + .write(true) + .open(path) + .expect("open sparse gguf"); + file.set_len(file_len).expect("ftruncate sparse gguf"); + file_len +} + #[test] fn mmap_matches_owned_parse_on_tiny_gguf() { let kv = [ @@ -81,6 +171,210 @@ fn mmap_matches_owned_parse_on_tiny_gguf() { ]); } +#[test] +fn mmap_directory_matches_rejects_different_tensor_names() { + let kv = [ + ("general.alignment", KvValue::U32(ALIGNMENT)), + ("general.architecture", KvValue::Str("olmoe")), + ]; + let payload = f32_vec_to_le_bytes(&[1.0, 2.0, 3.0, 4.0]); + let a_bytes = build_gguf( + &kv, + &[TensorSpec { + name: "token_embd.weight", + dims: vec![4], + ggml_type: GGML_F32, + payload: payload.clone(), + }], + ); + let b_bytes = build_gguf( + &kv, + &[TensorSpec { + name: "output.weight", + dims: vec![4], + ggml_type: GGML_F32, + payload, + }], + ); + let path = write_temp_gguf(&a_bytes); + let _guard = TempGuard(path.clone()); + let mapped = load_gguf_mmap(&path).expect("mmap a"); + let owned_b = parse_bytes(b_bytes, "mem://other".into()).expect("parse b"); + assert!( + !mapped.directory_matches(&owned_b), + "same count/arch but different tensor names must not match" + ); +} + +#[test] +fn mmap_dequant_matches_owned_for_k_quants() { + let page = os_page_size().max(4096); + let pad_elems = page; + let pad = vec![0u8; pad_elems * 4]; + let q8 = q8_0_ones_block(); + let q5 = q5_k_ones_block(); + let q6 = q6_k_minus_32_block(); + let iq3 = iq3_m_scaled_block(); + let kv = [ + ("general.alignment", KvValue::U32(ALIGNMENT)), + ("general.architecture", KvValue::Str("olmoe")), + ]; + let tensors = [ + TensorSpec { + name: "pad.weight", + dims: vec![pad_elems], + ggml_type: GGML_F32, + payload: pad, + }, + TensorSpec { + name: "q8.weight", + dims: vec![32], + ggml_type: GGML_Q8_0, + payload: q8.clone(), + }, + TensorSpec { + name: "q5.weight", + dims: vec![256], + ggml_type: GGML_Q5_K, + payload: q5.clone(), + }, + TensorSpec { + name: "q6.weight", + dims: vec![256], + ggml_type: GGML_Q6_K, + payload: q6.clone(), + }, + TensorSpec { + name: "iq3.weight", + dims: vec![256], + ggml_type: GGML_IQ3_M_BLOCK, + payload: iq3.clone(), + }, + ]; + let bytes = build_gguf(&kv, &tensors); + let path = write_temp_gguf(&bytes); + let _guard = TempGuard(path.clone()); + + let owned = load_gguf(&path).expect("load_gguf"); + let mapped = load_gguf_mmap(&path).expect("load_gguf_mmap"); + assert!( + mapped.directory_matches(&owned), + "mmap directory must match owned load" + ); + + let q8_t = mapped.tensor("q8.weight").expect("q8 tensor"); + assert!( + q8_t.absolute_offset >= page, + "q8 payload should start past the first OS page (offset {}, page {page})", + q8_t.absolute_offset + ); + let q8_pages = mapped + .tensor_page_aligned_bytes(q8_t) + .expect("q8 page range"); + assert!( + (q8_pages.pages.as_ptr() as usize).is_multiple_of(q8_pages.page_size), + "page-aligned range must start on an OS page" + ); + + let q8_map = mapped.tensor_bytes(q8_t).expect("q8 mmap"); + let q5_map = mapped + .tensor_bytes(mapped.tensor("q5.weight").unwrap()) + .expect("q5 mmap"); + let q6_map = mapped + .tensor_bytes(mapped.tensor("q6.weight").unwrap()) + .expect("q6 mmap"); + let iq3_map = mapped + .tensor_bytes(mapped.tensor("iq3.weight").unwrap()) + .expect("iq3 mmap"); + + let q8_owned = owned + .tensor_bytes(owned.tensor("q8.weight").unwrap()) + .expect("q8 owned"); + let q5_owned = owned + .tensor_bytes(owned.tensor("q5.weight").unwrap()) + .expect("q5 owned"); + let q6_owned = owned + .tensor_bytes(owned.tensor("q6.weight").unwrap()) + .expect("q6 owned"); + let iq3_owned = owned + .tensor_bytes(owned.tensor("iq3.weight").unwrap()) + .expect("iq3 owned"); + + assert_eq!(q8_map, q8_owned); + assert_eq!(q5_map, q5_owned); + assert_eq!(q6_map, q6_owned); + assert_eq!(iq3_map, iq3_owned); + assert_eq!(q8_pages.tensor_bytes(), q8_map); + + let q8_f = dequantize_q8_0(q8_map, &[32]).expect("q8 dequant"); + let q5_f = dequantize_q5_k(q5_map, &[256]).expect("q5 dequant"); + let q6_f = dequantize_q6_k(q6_map, &[256]).expect("q6 dequant"); + let iq3_f = dequantize_iq3_m(iq3_map, &[256]).expect("iq3 dequant"); + + assert_all(&[ + (q8_f.iter().all(|&v| v == 1.0), "q8 ones"), + (q5_f.iter().all(|&v| v == 1.0), "q5 ones"), + (q6_f.iter().all(|&v| v == -32.0), "q6 -32"), + (iq3_f[..16].iter().all(|&v| v == -3.0), "iq3 first scale"), + (iq3_f[16..].iter().all(|&v| v == 0.0), "iq3 rest zero"), + ( + dequantize_packed(DType::Q8_0, q8_map, &[32]).unwrap() == q8_f, + "packed q8", + ), + ( + dequantize_packed(DType::Q5_K, q5_map, &[256]).unwrap() == q5_f, + "packed q5", + ), + ( + dequantize_packed(DType::Q6_K, q6_map, &[256]).unwrap() == q6_f, + "packed q6", + ), + ( + dequantize_packed(DType::IQ3_M_BLOCK, iq3_map, &[256]).unwrap() == iq3_f, + "packed iq3", + ), + ]); +} + +#[test] +fn mmap_sparse_multi_gib_does_not_require_owned_read() { + let n_elements = SPARSE_PAYLOAD_BYTES / 4; + let path = { + let nanos = SystemTime::now() + .duration_since(UNIX_EPOCH) + .expect("clock") + .as_nanos(); + std::env::temp_dir().join(format!( + "engram-parser-mmap-sparse-{}-{nanos}.gguf", + process::id() + )) + }; + let file_len = write_sparse_f32_gguf(&path, n_elements); + let _guard = TempGuard(path.clone()); + + let allocated = fs::metadata(&path).expect("stat sparse").blocks() * 512; + if allocated >= 64 * 1024 * 1024 { + eprintln!( + "skip sparse multi-GiB mmap: filesystem allocated {allocated} bytes (not a hole)" + ); + return; + } + + let mapped = load_gguf_mmap(&path).expect("mmap sparse multi-GiB GGUF"); + assert_eq!(mapped.len() as u64, file_len, "mapped length"); + let tensor = mapped.tensor("sparse.weight").expect("sparse tensor"); + assert_eq!(tensor.byte_len as u64, SPARSE_PAYLOAD_BYTES); + assert_eq!(tensor.n_elements as u64, n_elements); + let prefix = mapped.tensor_bytes(tensor).expect("sparse payload view"); + assert_eq!(prefix.len() as u64, SPARSE_PAYLOAD_BYTES); + assert_eq!(&prefix[..16], &[0u8; 16], "sparse hole reads as zeros"); + let pages = mapped + .tensor_page_aligned_bytes(tensor) + .expect("sparse page range"); + assert_eq!(pages.tensor_bytes().len() as u64, SPARSE_PAYLOAD_BYTES); + assert!(pages.page_size > 0); +} + struct TempGuard(PathBuf); impl Drop for TempGuard {