todo: 创建 llama-nv

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

Author: YdrMaster

Cargo.toml

@@ -4,6 +4,7 @@ members = [
     "tensor",
     "models/llama/common",
     "models/llama/common-cpu",
+    "models/llama/nvidia-gpu",
     "test-utils",
 ]
 resolver = "2"
@@ -14,6 +15,6 @@ tensor.path = "tensor"
 causal-lm.path = "causal-lm"
 test-utils.path = "test-utils"
 
-ggus = { git = "https://github.com/YdrMaster/gguf", rev = "e64d758" }
+ggus = { git = "https://github.com/YdrMaster/gguf", rev = "c676bcc" }
 ndarray-layout = { git = "https://github.com/YdrMaster/ndarray-layout", rev = "5c6b969" }
 operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "64419f0", default-features = false }

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

@@ -142,15 +142,15 @@ struct Weights {
 
 impl WeightLoader for Weights {
     type Hardware = Cpu;
-    type Memory = &'static [u8];
+    type Memory<'s> = &'s [u8];
 
     #[inline]
     fn load_blk(
         &self,
         which: BlkWeight,
         iblk: usize,
         _queue: &QueueOf<Self::Hardware>,
-    ) -> Self::Memory {
+    ) -> Self::Memory<'_> {
         let blk = &self.blks[iblk];
         match which {
             BlkWeight::AttnNorm => blk.attn_norm,
@@ -163,12 +163,12 @@ impl WeightLoader for Weights {
     }
 
     #[inline]
-    fn output_norm(&self, _queue: &QueueOf<Self::Hardware>) -> Self::Memory {
+    fn output_norm(&self, _queue: &QueueOf<Self::Hardware>) -> Self::Memory<'_> {
         self.output_norm
     }
 
     #[inline]
-    fn output(&self, _queue: &QueueOf<Self::Hardware>) -> Self::Memory {
+    fn output(&self, _queue: &QueueOf<Self::Hardware>) -> Self::Memory<'_> {
         self.output
     }
 }

models/llama/common/src/compute.rs

@@ -54,17 +54,19 @@ pub enum BlkWeight {
 
 pub trait WeightLoader {
     type Hardware: Hardware;
-    type Memory: Deref<Target = [ByteOf<Self::Hardware>]>;
+    type Memory<'s>: Deref<Target = [ByteOf<Self::Hardware>]> + 's
+    where
+        Self: 's;
 
     fn load_blk(
         &self,
         which: BlkWeight,
         iblk: usize,
         queue: &QueueOf<Self::Hardware>,
-    ) -> Self::Memory;
+    ) -> Self::Memory<'_>;
 
-    fn output_norm(&self, queue: &QueueOf<Self::Hardware>) -> Self::Memory;
-    fn output(&self, queue: &QueueOf<Self::Hardware>) -> Self::Memory;
+    fn output_norm(&self, queue: &QueueOf<Self::Hardware>) -> Self::Memory<'_>;
+    fn output(&self, queue: &QueueOf<Self::Hardware>) -> Self::Memory<'_>;
 }
 
 pub struct LlamaWorker<Ops: Operators, W> {
@@ -544,60 +546,60 @@ impl LlamaMeta {
 
 impl<W: WeightLoader> WeightDecorator<W> {
     #[inline]
-    pub fn attn_norm(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory> {
+    pub fn attn_norm(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory<'_>> {
         combine(
             &self.attn_norm,
             self.weights.load_blk(BlkWeight::AttnNorm, iblk, queue),
         )
     }
 
     #[inline]
-    pub fn attn_qkv(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory> {
+    pub fn attn_qkv(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory<'_>> {
         combine(
             &self.attn_qkv,
             self.weights.load_blk(BlkWeight::AttnQKV, iblk, queue),
         )
     }
 
     #[inline]
-    pub fn attn_o(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory> {
+    pub fn attn_o(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory<'_>> {
         combine(
             &self.attn_o,
             self.weights.load_blk(BlkWeight::AttnO, iblk, queue),
         )
     }
 
     #[inline]
-    pub fn ffn_norm(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory> {
+    pub fn ffn_norm(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory<'_>> {
         combine(
             &self.ffn_norm,
             self.weights.load_blk(BlkWeight::FfnNorm, iblk, queue),
         )
     }
 
     #[inline]
-    pub fn ffn_gate_up(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory> {
+    pub fn ffn_gate_up(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory<'_>> {
         combine(
             &self.ffn_gate_up,
             self.weights.load_blk(BlkWeight::FfnGateUp, iblk, queue),
         )
     }
 
     #[inline]
-    pub fn ffn_down(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory> {
+    pub fn ffn_down(&self, iblk: usize, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory<'_>> {
         combine(
             &self.ffn_down,
             self.weights.load_blk(BlkWeight::FfnDown, iblk, queue),
         )
     }
 
     #[inline]
-    pub fn output_norm(&self, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory> {
+    pub fn output_norm(&self, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory<'_>> {
         combine(&self.output_norm, self.weights.output_norm(queue))
     }
 
     #[inline]
-    pub fn output(&self, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory> {
+    pub fn output(&self, queue: &QueueOf<W::Hardware>) -> Tensor<W::Memory<'_>> {
         combine(&self.output, self.weights.output(queue))
     }
 }

models/llama/nvidia-gpu/Cargo.toml

@@ -0,0 +1,15 @@
+[package]
+name = "llama-nv"
+version = "0.0.0"
+edition = "2021"
+authors = ["YdrMaster <ydrml@hotmail.com>"]
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+llama.path = "../common"
+operators = { workspace = true, features = ["nvidia-gpu"] }
+
+[dev-dependencies]
+test-utils.workspace = true
+gguf.workspace = true

models/llama/nvidia-gpu/src/lib.rs

@@ -0,0 +1,76 @@
+use llama::{ext::Mmap, LlamaStorage, Tensor, WeightLoader};
+use operators::{
+    all_reduce::NonAllReduce,
+    cuda::{memcpy_d2h, DevByte},
+    nvidia_gpu::Gpu,
+    ByteOf,
+};
+use std::ops::Deref;
+
+pub struct Llama {}
+
+impl Llama {
+    pub fn new(_storage: Box<[Mmap]>, _model: LlamaStorage<&'static [u8]>) -> Self {
+        Self {}
+    }
+
+    pub fn infer(&mut self, input: &[u32], cache: &mut [u8], pos: usize) -> u32 {
+        todo!()
+    }
+}
+
+struct Operators;
+
+macro_rules! op {
+    ($name:ident) => {
+        operators::$name::nvidia_gpu::Operator
+    };
+}
+
+impl llama::Operators for Operators {
+    type Hardware = Gpu;
+    type TopoNode = Gpu;
+    type RmsNorm = op!(rms_norm);
+    type MatMul = op!(mat_mul);
+    type Rope = op!(rope);
+    type AttnKVCached = op!(attention_kv_cached);
+    type Mlp = op!(mlp);
+    type Rearrange = op!(rearrange);
+    type AllReduce = NonAllReduce<Gpu>;
+
+    fn debug<T>(tensor: &Tensor<T>)
+    where
+        T: Deref<Target = [ByteOf<Self::Hardware>]>,
+    {
+        let tensor = tensor.as_ref().map(|mem| {
+            let mut buf = vec![0u8; mem.len()];
+            memcpy_d2h(&mut buf, mem);
+            buf
+        });
+        println!("{tensor}");
+    }
+}
+
+struct Weights {}
+
+impl WeightLoader for Weights {
+    type Hardware = Gpu;
+    type Memory<'s> = &'s [DevByte];
+
+    fn load_blk(
+        &self,
+        which: llama::BlkWeight,
+        iblk: usize,
+        queue: &operators::QueueOf<Self::Hardware>,
+    ) -> Self::Memory<'_> {
+        todo!()
+    }
+
+    fn output_norm(&self, queue: &operators::QueueOf<Self::Hardware>) -> Self::Memory<'_> {
+        todo!()
+    }
+
+    fn output(&self, queue: &operators::QueueOf<Self::Hardware>) -> Self::Memory<'_> {
+        todo!()
+    }
+}