build: 梳理模块间依赖关系

Signed-off-by: YdrMaster <[email protected]>

Cargo.toml

@@ -15,8 +15,5 @@ causal-lm.path = "causal-lm"
 test-utils.path = "test-utils"
 
 ggus = { git = "https://github.com/YdrMaster/gguf", rev = "e64d758" }
-ggml-quants = { git = "https://github.com/YdrMaster/gguf", rev = "e64d758" }
 ndarray-layout = { git = "https://github.com/YdrMaster/ndarray-layout", rev = "5c6b969" }
 operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "a9ddac5", default-features = false }
-
-memmap2 = "0.9"

gguf/Cargo.toml

@@ -6,10 +6,11 @@ authors = ["YdrMaster <[email protected]>"]
 
 [dependencies]
 ggus.workspace = true
-memmap2.workspace = true
 minijinja = { version = "2.3", default-features = false, features = ["loader"] }
 serde = { version = "1.0", features = ["derive"] }
 tokeneer = "0.0"
+memmap2 = "0.9"
+
 
 [dev-dependencies]
 test-utils.workspace = true

gguf/src/chat_template.rs

@@ -1,4 +1,4 @@
-use super::{GGufModel, Tokenizer};
+use super::{tokenizer::Tokenizer, GGufModel};
 use ggus::{GGufMetaError, GGufMetaMapExt};
 use minijinja::Environment;
 use serde::Serialize;

gguf/src/lib.rs

@@ -10,7 +10,9 @@ use memmap2::Mmap;
 use std::{collections::HashMap, fmt::Debug, fs::File, path::Path};
 
 pub use ggus::{ggml_quants, GGufMetaError, GGufMetaMapExt};
-pub use tokenizer::Tokenizer;
+pub mod ext {
+    pub use memmap2::Mmap;
+}
 
 /// 从指定文件的路径出发，映射所有分片文件。
 pub fn map_files(path: impl AsRef<Path>) -> Box<[Mmap]> {

models/llama/common-cpu/Cargo.toml

@@ -9,9 +9,6 @@ authors = ["YdrMaster <[email protected]>"]
 [dependencies]
 llama.path = "../common"
 operators = { workspace = true, features = ["common-cpu"] }
-memmap2.workspace = true
-tensor.workspace = true
-half = "2.4"
 
 [dev-dependencies]
 test-utils.workspace = true

models/llama/common-cpu/src/lib.rs

@@ -1,16 +1,14 @@
-use half::f16;
 use llama::{
-    primitive, BlkWeight, LlamaArgs, LlamaBlkStorage, LlamaBlks, LlamaMeta, LlamaRequest,
-    LlamaStorage, RandomSample, WeightLoader,
+    ext::{f16, primitive, Mmap},
+    BlkWeight, LlamaArgs, LlamaBlkStorage, LlamaBlks, LlamaMeta, LlamaRequest, LlamaStorage,
+    RandomSample, Tensor, WeightLoader,
 };
-use memmap2::Mmap;
 use operators::{
     common_cpu::{Cpu, ThisThread},
     random_sample::{common_cpu::Operator as CpuOp, KVPair, SampleArgs},
     ByteOf, QueueOf,
 };
 use std::{ops::Deref, slice::from_raw_parts_mut};
-use tensor::{ArrayLayout, BigEndian, Tensor};
 
 pub struct Llama {
     _storage: Box<[Mmap]>,
@@ -93,13 +91,9 @@ impl Llama {
             .unwrap();
 
         let mut pair = KVPair::new(0, f16::ZERO);
-        let mut pairs = unsafe {
-            Tensor::from_raw_parts(
-                KVPair::<()>::LAYOUT,
-                ArrayLayout::new_contiguous(&[1], BigEndian, size_of_val(&pair)),
-                from_raw_parts_mut(&mut pair as *mut _ as *mut u8, size_of_val(&pair)),
-            )
-        };
+        let mut pairs = Tensor::kv_pair_vec(1, |_| unsafe {
+            from_raw_parts_mut(&mut pair as *mut _ as *mut u8, size_of_val(&pair))
+        });
 
         self.sample
             .launch(
@@ -248,7 +242,7 @@ fn test_infer() {
 
     fn print_time(name: &str, time: Duration, n: usize) {
         println!(
-            "{name} : {time:?} for {n} tokens, avg: {:?} per token",
+            "{name}: {time:?} for {n} tokens, avg: {:?} per token",
             time.div_f64(n as _)
         );
     }

models/llama/common/src/lib.rs

@@ -3,18 +3,25 @@ mod compute;
 mod random_sample;
 mod storage;
 
-use tensor::Tensor;
-
 pub use args::{Args as LlamaArgs, Request as LlamaRequest};
 pub use compute::{BlkWeight, LlamaBlks, Operators, WeightLoader};
-pub use gguf::ggml_quants::digit_layout::{types as primitive, DigitLayout};
 pub use random_sample::RandomSample;
 pub use storage::{BlkStorage as LlamaBlkStorage, Storage as LlamaStorage};
+pub use tensor::Tensor;
+pub mod ext {
+    pub use gguf::{
+        ext::Mmap,
+        ggml_quants::{
+            digit_layout::{types as primitive, DigitLayout},
+            f16, types as quant,
+        },
+    };
+}
 
 #[derive(Clone, Debug)]
 pub struct LlamaMeta {
-    pub dt_norm: DigitLayout,
-    pub dt_mat: DigitLayout,
+    pub dt_norm: ext::DigitLayout,
+    pub dt_mat: ext::DigitLayout,
     pub nblk: usize,
     pub nh: usize,
     pub nkvh: usize,

tensor/Cargo.toml

@@ -5,7 +5,6 @@ edition = "2021"
 authors = ["YdrMaster <[email protected]>"]
 
 [dependencies]
-ggml-quants.workspace = true
-ndarray-layout.workspace = true
+ggus.workspace = true
 operators.workspace = true
-half = "2.4"
+ndarray-layout.workspace = true

tensor/src/fmt.rs

@@ -1,6 +1,5 @@
 use super::Tensor;
-use ggml_quants::digit_layout::types as primitive;
-use half::{bf16, f16};
+use ggus::ggml_quants::{bf16, digit_layout::types as primitive, f16};
 use std::{fmt, ops::Deref};
 
 pub trait DataFmt {

tensor/src/lib.rs

@@ -1,14 +1,15 @@
 mod fmt;
 mod split;
 
-use ggml_quants::{digit_layout::DigitLayout, DataBlock};
-use operators::TensorLayout;
+use ggus::ggml_quants::{digit_layout::DigitLayout, DataBlock};
+use ndarray_layout::{ArrayLayout, Endian::BigEndian};
+use operators::{random_sample::KVPair, TensorLayout};
 use std::{
+    alloc::Layout,
     ops::{Deref, DerefMut, Range},
     slice::from_raw_parts,
 };
 
-pub use ndarray_layout::{ArrayLayout, Endian::BigEndian};
 pub use split::{LocalSplitable, Splitable};
 
 #[derive(Clone)]
@@ -19,19 +20,6 @@ pub struct Tensor<T> {
 }
 
 impl<T> Tensor<T> {
-    #[inline]
-    pub const unsafe fn from_raw_parts(
-        element: DigitLayout,
-        layout: ArrayLayout<5>,
-        physical: T,
-    ) -> Self {
-        Self {
-            element,
-            layout,
-            physical,
-        }
-    }
-
     pub fn new(element: DigitLayout, shape: &[usize], physical: T) -> Self {
         Self {
             element,
@@ -107,11 +95,6 @@ impl<T> Tensor<T> {
             physical: f(self.physical),
         }
     }
-
-    #[inline]
-    pub fn layout(&self) -> TensorLayout {
-        TensorLayout::new(self.element, self.layout.shape(), self.layout.strides())
-    }
 }
 
 impl<T, B> Tensor<T>
@@ -153,6 +136,22 @@ where
     }
 }
 
+/// operators
+impl<T> Tensor<T> {
+    pub fn kv_pair_vec(n: usize, f: impl FnOnce(usize) -> T) -> Self {
+        Self {
+            element: KVPair::<()>::LAYOUT,
+            layout: ArrayLayout::new_contiguous(&[n], BigEndian, size_of::<KVPair>()),
+            physical: f(Layout::array::<KVPair>(n).unwrap().size()),
+        }
+    }
+
+    #[inline]
+    pub fn layout(&self) -> TensorLayout {
+        TensorLayout::new(self.element, self.layout.shape(), self.layout.strides())
+    }
+}
+
 /// transform
 impl<T> Tensor<T> {
     #[inline]

test-utils/Cargo.toml

@@ -6,4 +6,3 @@ authors = ["YdrMaster <[email protected]>"]
 
 [dependencies]
 gguf.workspace = true
-memmap2.workspace = true

test-utils/src/lib.rs

@@ -1,5 +1,4 @@
-use gguf::map_files;
-use memmap2::Mmap;
+use gguf::{ext::Mmap, map_files};
 use std::path::Path;
 
 pub fn map_gguf_files() -> Option<Box<[Mmap]>> {