fix: 实现 cpu 推理

Signed-off-by: YdrMaster <ydrml@hotmail.com>

Author: YdrMaster

Cargo.lock

@@ -17,6 +17,6 @@ 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 = "9b5c6b9", default-features = false }
+operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "656e7f7", default-features = false }
 
 memmap2 = "0.9"

Cargo.toml

@@ -1,13 +1,15 @@
 mod chat_template;
 mod tokenizer;
 
+use ggml_quants::digit_layout::DigitLayout;
 use ggus::{
-    ggml_quants::digit_layout::DigitLayout, GGuf, GGufError, GGufFileName, GGufMetaDataValueType,
-    GGufMetaKV, GGufMetaMap, GGufReadError, GENERAL_ALIGNMENT,
+    GGuf, GGufError, GGufFileName, GGufMetaDataValueType, GGufMetaKV, GGufMetaMap, GGufReadError,
+    GENERAL_ALIGNMENT,
 };
 use memmap2::Mmap;
 use std::{collections::HashMap, fmt::Debug, fs::File, path::Path};
 
+pub use ggus::{ggml_quants, GGufMetaError, GGufMetaMapExt};
 pub use tokenizer::Tokenizer;
 
 /// 从指定文件的路径出发，映射所有分片文件。

gguf/src/lib.rs

@@ -7,9 +7,9 @@ use memmap2::Mmap;
 use operators::{
     common_cpu::{Cpu, ThisThread},
     random_sample::{common_cpu::Operator as CpuOp, KVPair, SampleArgs},
-    QueueOf,
+    ByteOf, QueueOf,
 };
-use std::slice::from_raw_parts_mut;
+use std::{ops::Deref, slice::from_raw_parts_mut};
 use tensor::{ArrayLayout, BigEndian, Tensor};
 
 pub struct Llama {
@@ -62,7 +62,7 @@ impl Llama {
         let mut embd_buf = vec![0u8; embd.shape().iter().product::<usize>() * ele];
         let mut logits_buf = vec![0u8; logits.shape().iter().product::<usize>() * ele];
 
-        let d = embd.shape()[1];
+        let d = embd.shape()[1] * ele;
         for (i, &tok) in input.iter().enumerate() {
             embd_buf[i * d..][..d].copy_from_slice(&self.token_embed[tok as usize * d..][..d]);
         }
@@ -132,6 +132,13 @@ impl llama::Operators for Operators {
     type AttnKVCached = op!(attention_kv_cached);
     type Mlp = op!(mlp);
     type Rearrange = op!(rearrange);
+
+    fn debug<T>(tensor: &Tensor<T>)
+    where
+        T: Deref<Target = [ByteOf<Self::Hardware>]>,
+    {
+        println!("{tensor}");
+    }
 }
 
 struct Weights {
@@ -174,14 +181,19 @@ impl WeightLoader for Weights {
 }
 
 #[test]
-fn test_load() {
-    use gguf::GGufModel;
-    use std::{io::Write, slice::from_raw_parts};
+fn test_infer() {
+    use gguf::{GGufMetaMapExt, GGufModel};
+    use std::{
+        io::Write,
+        slice::from_raw_parts,
+        time::{Duration, Instant},
+    };
 
     let Some(shards) = test_utils::map_gguf_files() else {
         return;
     };
     let gguf = GGufModel::read(shards.iter().map(|s| &**s));
+    let eos = gguf.tokenizer_ggml_eos_token_id().unwrap();
     let tokenizer = gguf.tokenizer();
     let llama =
         LlamaStorage::from_gguf(&gguf).map(&mut |s| unsafe { from_raw_parts(s.as_ptr(), s.len()) });
@@ -194,14 +206,50 @@ fn test_load() {
     let mut cache_buf = vec![0u8; cache.shape().iter().product::<usize>() * size_of::<f16>()];
 
     let mut prompt = "Once upon a time,".to_string();
+
+    print!("{prompt}");
+    std::io::stdout().flush().unwrap();
+
     let mut tokens = tokenizer.encode(&prompt);
-    while !tokens.contains(&2) {
-        let next = llama.infer(&tokens, &mut cache_buf, 0);
-        tokens = vec![next];
+    let num_prompt_tokens = tokens.len();
+
+    let mut prefill = Duration::ZERO;
+    let mut decode = Duration::ZERO;
+
+    let mut pos = 0;
+    loop {
+        let time = Instant::now();
+        let next = llama.infer(&tokens, &mut cache_buf, pos);
+        let time = time.elapsed();
+
+        if prefill.is_zero() {
+            prefill = time;
+        } else {
+            decode += time;
+        }
+
+        pos += tokens.len();
+        if next == eos {
+            break;
+        }
 
         let piece = tokenizer.decode(next);
         print!("{piece}");
         std::io::stdout().flush().unwrap();
         prompt.push_str(&piece);
+        tokens = vec![next];
+    }
+
+    println!();
+    println!();
+    print_time("total", prefill + decode, pos);
+    print_time("prefill", prefill, num_prompt_tokens);
+    print_time("decode", decode, pos - num_prompt_tokens);
+
+    fn print_time(name: &str, time: Duration, n: usize) {
+        println!(
+            "{name} : {time:?} for {n} tokens, avg: {:?} per token",
+            time.div_f64(n as _)
+        );
     }
 }

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

@@ -7,7 +7,6 @@ authors = ["YdrMaster <ydrml@hotmail.com>"]
 # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
 
 [dependencies]
-ggus.workspace = true
 gguf.workspace = true
 operators.workspace = true
 tensor.workspace = true

models/llama/common/Cargo.toml

@@ -1,5 +1,5 @@
 use super::{args::Args, LlamaMeta};
-use ggus::ggml_quants::digit_layout::types as ty;
+use gguf::ggml_quants::digit_layout::types as ty;
 use itertools::izip;
 use operators::{
     attention_kv_cached::AttnKVCached,
@@ -21,6 +21,10 @@ pub trait Operators {
     type AttnKVCached: AttnKVCached<Self::Hardware>;
     type Mlp: Mlp<Self::Hardware>;
     type Rearrange: Rearrange<Self::Hardware>;
+
+    fn debug<T>(tensor: &Tensor<T>)
+    where
+        T: Deref<Target = [ByteOf<Self::Hardware>]>;
 }
 
 pub enum BlkWeight {
@@ -255,8 +259,8 @@ where
         let x_ = unsafe { x.map_slice_static() };
         self.rms_norm(&mut x, &x_, &w, workspace, queue_alloc)?;
 
-        let lm_head = self.weights.output(queue);
-        self.mat_mul(&mut logits, 0., &x, &lm_head, 1., workspace, queue_alloc)
+        let output = self.weights.output(queue);
+        self.mat_mul(&mut logits, 0., &x, &output, 1., workspace, queue_alloc)
     }
 }
 

models/llama/common/src/compute.rs

@@ -3,12 +3,11 @@ mod compute;
 mod random_sample;
 mod storage;
 
-use ggus::ggml_quants::digit_layout::DigitLayout;
 use tensor::Tensor;
 
 pub use args::{Args as LlamaArgs, Request as LlamaRequest};
 pub use compute::{BlkWeight, LlamaBlks, Operators, WeightLoader};
-pub use ggus::ggml_quants::digit_layout::types as primitive;
+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};
 

models/llama/common/src/lib.rs

@@ -1,6 +1,5 @@
 use crate::LlamaMeta;
-use gguf::GGufModel;
-use ggus::{GGufMetaError, GGufMetaMapExt};
+use gguf::{GGufMetaError, GGufMetaMapExt, GGufModel};
 
 #[derive(Clone)]
 pub struct Storage<T> {