feat(transformer): 抽取 transformer 平台无关部分到 crate

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

Cargo.toml

@@ -4,6 +4,7 @@ members = [
     "tensor",
     "model-parameters",
     "tokenizer",
+    "transformer",
     "transformer-cpu",
     "transformer-nvidia",
     "xtask",

transformer-cpu/Cargo.toml

@@ -10,6 +10,7 @@ authors = ["YdrMaster <[email protected]>"]
 common = { path = "../common" }
 tensor = { path = "../tensor" }
 model-parameters = { path = "../model-parameters" }
+transformer = { path = "../transformer" }
 gemm = "0.17"
 
 [dev-dependencies]

transformer-cpu/src/kernel/gather.rs

@@ -3,7 +3,7 @@ use crate::Request;
 use std::ops::{Deref, DerefMut};
 use tensor::{udim, Tensor};
 
-pub fn gather<T, U>(mut x: Tensor<T>, table: &Tensor<U>, requests: &[Request])
+pub fn gather<T, U, X>(mut x: Tensor<T>, table: &Tensor<U>, requests: &[Request<X>])
 where
     T: DerefMut<Target = [u8]>,
     U: Deref<Target = [u8]>,

transformer-cpu/src/lib.rs

@@ -1,55 +1,20 @@
-mod cache;
 mod kernel;
 mod storage;
 
-use common::{upos, utok};
 use gemm::f16;
 use kernel::{gather, mat_mul, rms_norm, rms_norm_inplace, rotary_embedding, softmax, swiglu};
 use model_parameters::{Llama2, Memory};
 use storage::Storage;
 use tensor::{reslice, reslice_mut, slice, udim, DataType, Tensor};
 
-pub use cache::LayerCache;
+pub type LayerCache = transformer::LayerCache<Storage>;
+pub use transformer::{Prompt, Request};
 pub extern crate model_parameters;
 
 pub struct Transformer {
     model: Box<dyn Llama2>,
 }
 
-pub struct Request<'a> {
-    pub prompt: Prompt<'a>,
-    pub cache: &'a mut [LayerCache],
-    pub pos: upos,
-}
-
-pub enum Prompt<'a> {
-    Prefill(&'a [utok]),
-    Decode(utok),
-}
-
-impl Request<'_> {
-    #[inline]
-    pub const fn tokens(&self) -> &[utok] {
-        match &self.prompt {
-            Prompt::Prefill(tokens) => tokens,
-            Prompt::Decode(token) => std::slice::from_ref(&token),
-        }
-    }
-
-    #[inline]
-    pub const fn seq_len(&self) -> udim {
-        match self.prompt {
-            Prompt::Prefill(tokens) => tokens.len() as _,
-            Prompt::Decode(_) => 1,
-        }
-    }
-
-    #[inline]
-    pub const fn att_len(&self) -> udim {
-        self.pos + self.seq_len()
-    }
-}
-
 impl Transformer {
     #[inline]
     pub fn new(model: Box<dyn Llama2>) -> Self {
@@ -63,15 +28,15 @@ impl Transformer {
 
     #[inline]
     pub fn new_cache(&self) -> Vec<LayerCache> {
-        LayerCache::new_layers(&*self.model)
+        LayerCache::new_layers(&*self.model, tensor)
     }
 
     #[inline]
     pub fn max_seq_len(&self) -> usize {
         self.model.max_position_embeddings()
     }
 
-    pub fn decode(&mut self, mut requests: Vec<Request>) -> Vec<f16> {
+    pub fn decode(&mut self, mut requests: Vec<Request<Storage>>) -> Vec<f16> {
         use std::cmp::Ordering::*;
         requests.sort_unstable_by(|a, b| match a.prompt {
             Prompt::Prefill(_) => match b.prompt {

transformer-nvidia/Cargo.toml

@@ -10,6 +10,7 @@ authors = ["YdrMaster <[email protected]>"]
 common = { path = "../common" }
 tensor = { path = "../tensor" }
 model-parameters = { path = "../model-parameters" }
+transformer = { path = "../transformer" }
 cuda = { git = "https://github.com/YdrMaster/cuda-bench" }
 cublas = { git = "https://github.com/YdrMaster/cuda-bench" }
 half.workspace = true

transformer-nvidia/src/lib.rs

@@ -1,9 +1,8 @@
 #![cfg(detected_cuda)]
 
-mod cache;
 mod kernel;
+mod page_locked_memory;
 mod parameters;
-mod storage;
 
 use ::half::f16;
 use common::{upos, utok};
@@ -15,8 +14,8 @@ use parameters::{LayersParameters, ModelParameters};
 use std::ptr::null_mut;
 use tensor::{slice, udim, DataType, Tensor};
 
-pub use cache::LayerCache;
-pub use storage::PageLockedMemory;
+pub type LayerCache<'a> = transformer::LayerCache<LocalDevBlob<'a>>;
+pub use page_locked_memory::PageLockedMemory;
 pub extern crate cuda;
 pub extern crate model_parameters;
 
@@ -73,13 +72,13 @@ impl<'a> Transformer<'a> {
 
     #[inline]
     pub fn new_cache<'b>(&self, stream: &'b Stream) -> Vec<LayerCache<'b>> {
-        LayerCache::new_layers(self.host, stream)
+        LayerCache::new_layers(self.host, |dt, shape| tensor(dt, shape, stream))
     }
 
     pub fn update<'b>(
         &mut self,
         tokens: &[utok],
-        cache: &[LayerCache],
+        cache: &mut [LayerCache],
         pos: upos,
         compute: &Stream,
         transfer: &'b Stream,
@@ -128,7 +127,7 @@ impl<'a> Transformer<'a> {
 
         cublas!(cublasSetStream_v2(self.cublas, compute.as_raw() as _));
         compute.wait_for(&e_alloc);
-        for (layer, cache) in cache.iter().enumerate() {
+        for (layer, cache) in cache.iter_mut().enumerate() {
             self.layers.load(layer, self.host, transfer);
             let params = self.layers.sync(layer, compute);
 
@@ -248,7 +247,7 @@ impl<'a> Transformer<'a> {
     pub fn decode(
         &mut self,
         token: utok,
-        cache: &[LayerCache],
+        cache: &mut [LayerCache],
         pos: upos,
         compute: &Stream,
         transfer: &Stream,

transformer/Cargo.toml

@@ -0,0 +1,12 @@
+[package]
+name = "transformer"
+version = "0.0.0"
+edition = "2021"
+authors = ["YdrMaster <[email protected]>"]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+common = { path = "../common" }
+tensor = { path = "../tensor" }
+model-parameters = { path = "../model-parameters" }

transformer-cpu/src/cache.rs → transformer/src/cache.rs

@@ -1,21 +1,27 @@
-use crate::{tensor, Storage};
-use model_parameters::Llama2;
-use tensor::{udim, Tensor};
+use model_parameters::Llama2;
+use tensor::{udim, DataType, Tensor};
 
-pub struct LayerCache {
+/// KV cache for one layer.
+pub struct LayerCache<Storage> {
     /// Key cache, shape = `num_kv_head x max_seq_len x head_dim`.
     k: Tensor<Storage>,
     /// Value cache, shape = `num_kv_head x max_seq_len x head_dim`.
     v: Tensor<Storage>,
 }
 
-impl LayerCache {
-    pub fn new_layers(model: &dyn Llama2) -> Vec<Self> {
-        let dt = model.data_type();
+impl<Storage> LayerCache<Storage> {
+    /// Alloc KV Cache for all layers.
+    pub fn new_layers(
+        model: &dyn Llama2,
+        tensor: impl Fn(DataType, &[udim]) -> Tensor<Storage>,
+    ) -> Vec<Self> {
         let nkvh = model.num_key_value_heads() as udim;
-        let hd = (model.hidden_size() / model.num_attention_heads()) as udim;
         let max_seq_len = model.max_position_embeddings() as udim;
-        let shape = &[nkvh, max_seq_len, hd];
+        let dh = (model.hidden_size() / model.num_attention_heads()) as udim;
+
+        let dt = model.data_type();
+        let shape = &[nkvh, max_seq_len, dh];
+
         (0..model.num_hidden_layers())
             .map(|_| Self {
                 k: tensor(dt, shape),
@@ -24,6 +30,7 @@ impl LayerCache {
             .collect()
     }
 
+    /// Get mutable references to the key and value cache.
     #[inline]
     pub fn get(&mut self) -> (&mut Tensor<Storage>, &mut Tensor<Storage>) {
         (&mut self.k, &mut self.v)

transformer/src/lib.rs

@@ -0,0 +1,54 @@
+//! Common code for transformers.
+
+#![deny(warnings, missing_docs)]
+
+mod cache;
+
+use common::{upos, utok};
+use tensor::udim;
+
+pub use cache::LayerCache;
+
+/// A request to decode a sequence.
+pub struct Request<'a, Storage> {
+    /// Prompt of this request.
+    pub prompt: Prompt<'a>,
+    /// Context cache of this request.
+    pub cache: &'a mut [LayerCache<Storage>],
+    /// Position of `prompt` in context.
+    pub pos: upos,
+}
+
+/// User prompt in transformer inference once.
+pub enum Prompt<'a> {
+    /// Prefill the sequence with tokens.
+    Prefill(&'a [utok]),
+    /// Decode the next token.
+    Decode(utok),
+}
+
+impl<S> Request<'_, S> {
+    /// Tokens in the prompt.
+    #[inline]
+    pub const fn tokens(&self) -> &[utok] {
+        match &self.prompt {
+            Prompt::Prefill(tokens) => tokens,
+            Prompt::Decode(token) => std::slice::from_ref(&token),
+        }
+    }
+
+    /// Length of tokens in the prompt.
+    #[inline]
+    pub const fn seq_len(&self) -> udim {
+        match self.prompt {
+            Prompt::Prefill(tokens) => tokens.len() as _,
+            Prompt::Decode(_) => 1,
+        }
+    }
+
+    /// Length of tokens in attention computation.
+    #[inline]
+    pub const fn att_len(&self) -> udim {
+        self.pos + self.seq_len()
+    }
+}

xtask/src/generate.rs

@@ -253,7 +253,7 @@ fn on_nvidia_gpu(
         let host = Memory::load_safetensors(config, host, false).unwrap();
         let eos = host.eos_token_id();
         let mut transformer = Transformer::new(&host, preload_layers, &transfer);
-        let kv_cache = transformer.new_cache(&compute);
+        let mut kv_cache = transformer.new_cache(&compute);
         info!("build model host: {:?}", time.elapsed());
 
         let step = step.min(host.max_position_embeddings());
@@ -264,7 +264,7 @@ fn on_nvidia_gpu(
         let time = Instant::now();
         let (last, tokens) = prompt_tokens.split_last().expect("prompt is empty");
         if !tokens.is_empty() {
-            transformer.update(tokens, &kv_cache, 0, &compute, &transfer);
+            transformer.update(tokens, &mut kv_cache, 0, &compute, &transfer);
         }
         info!("prefill transformer ... {:?}", time.elapsed());
 
@@ -274,7 +274,7 @@ fn on_nvidia_gpu(
         let mut pos = tokens.len();
         let time = Instant::now();
         while pos < step {
-            let logits = transformer.decode(token, &kv_cache, pos as _, &compute, &transfer);
+            let logits = transformer.decode(token, &mut kv_cache, pos as _, &compute, &transfer);
             token = argmax(logits);
 
             print!("{}", tokenizer.decode(token).replace('▁', " "));