Merge pull request #13 from InfiniTensor/dev-operators

move random search to operators-rs

Author: YdrMaster

.github/workflows/build.yml

@@ -26,6 +26,11 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
 
+      - name: cuda-toolkit
+        uses: Jimver/cuda-toolkit@v0.2.16
+        with:
+          method: 'network'
+
       - name: Check format
         run: cargo fmt --check
 

Cargo.lock

@@ -4,7 +4,6 @@ members = [
     "tensor",
     "tokenizer",
     "causal-lm",
-    "sample",
     "service",
     "web-api",
     "xtask",
@@ -35,7 +34,6 @@ tokio = { version = "1.38", features = ["rt-multi-thread", "sync"] }
 digit-layout = "0.0"
 build-script-cfg = "0.0"
 
-operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "04e71d5", default-features = false }
-nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "343b0e0" }
-search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "343b0e0" }
+operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "5a88159", default-features = false }
+search-cuda-tools = { git = "https://github.com/YdrMaster/cuda-driver", rev = "fb088b6" }
 search-neuware-tools = "0.0"

Cargo.toml

@@ -7,6 +7,8 @@
 ![GitHub contributors](https://img.shields.io/github/contributors/InfiniTensor/transformer-rs)
 ![GitHub commit activity](https://img.shields.io/github/commit-activity/m/InfiniTensor/transformer-rs)
 
+[**使用指南**](/docs/user-guide/doc.md)
+
 从 [YdrMaster/llama2.rs](https://github.com/YdrMaster/llama2.rs) 发展来的手写 transformer 模型项目。
 
 ## 使用

README.md

@@ -9,5 +9,5 @@ authors = ["YdrMaster <ydrml@hotmail.com>"]
 [dependencies]
 common = { path = "../common" }
 tensor = { path = "../tensor" }
-sample = { path = "../sample" }
 digit-layout.workspace = true
+operators.workspace = true

causal-lm/Cargo.toml

@@ -10,8 +10,8 @@ use std::path::Path;
 use tensor::{udim, Tensor};
 
 pub use decoding::DecodingMeta;
+pub use operators::random_sample::SampleArgs;
 pub use query_context::QueryContext;
-pub use sample::SampleArgs;
 
 /// 从文件系统加载的模型。
 pub trait Model: Sized {
@@ -142,7 +142,7 @@ where
 
         let args = [SampleMeta {
             num_decode: 1,
-            args: SampleArgs::default(),
+            args: SampleArgs::ARG_MAX,
         }];
         let tokens = CausalLM::sample(&model, args, logits);
 

causal-lm/src/lib.rs

@@ -7,12 +7,18 @@ macro_rules! slice {
 
 mod gather;
 
-use common::utok;
+use common::{f16, utok};
 use common_devices::{Operators, SliceOn};
+use digit_layout::types::F16;
 use operators::{
-    fuesd_softmax::common_cpu as softmax, mat_mul::common_cpu as mat_mul,
-    reform::common_cpu as reform, rms_norm::common_cpu as rms_norm, rope::common_cpu as rope,
-    swiglu::common_cpu as swiglu, Operator, QueueOf,
+    fuesd_softmax::common_cpu as softmax,
+    mat_mul::common_cpu as mat_mul,
+    random_sample::{common_cpu as random_sample, Args, KVPair, SampleArgs},
+    reform::common_cpu as reform,
+    rms_norm::common_cpu as rms_norm,
+    rope::common_cpu as rope,
+    swiglu::common_cpu as swiglu,
+    Operator, QueueOf,
 };
 use std::ops::{Deref, DerefMut};
 use tensor::Tensor;
@@ -29,6 +35,23 @@ pub struct CpuKernels {
     rope: rope::Operator,
     softmax: softmax::Operator,
     swiglu: swiglu::Operator,
+    sample: random_sample::Operator,
+}
+
+impl CpuKernels {
+    pub fn sample(&self, temperature: f32, top_p: f32, top_k: usize, logits: &[f16]) -> utok {
+        let mut kv_pair = KVPair::new(0, f16::ZERO);
+        let mut args = Args::<Cpu>::new(F16, logits.len());
+        args.kv_pair_base = &mut kv_pair as *mut _ as _;
+        args.data_base = logits.as_ptr() as _;
+        args.detail = SampleArgs {
+            temperature,
+            top_p,
+            top_k,
+        };
+        self.sample.launch(&args, &ThisThread).unwrap();
+        kv_pair.idx() as _
+    }
 }
 
 impl Default for CpuKernels {
@@ -40,6 +63,7 @@ impl Default for CpuKernels {
             rope: rope::Operator::new(&Cpu),
             softmax: softmax::Operator::new(&Cpu),
             swiglu: swiglu::Operator::new(&Cpu),
+            sample: random_sample::Operator::new(&Cpu),
         }
     }
 }

devices/common-cpu/src/lib.rs

@@ -10,12 +10,10 @@ authors = ["YdrMaster <ydrml@hotmail.com>"]
 common = { path = "../../common" }
 common-devices = { path = "../common" }
 tensor = { path = "../../tensor" }
-sample = { path = "../../sample" }
 rand = "0.8"
 operators = { workspace = true, features = ["nvidia-gpu"] }
 digit-layout.workspace = true
 
 [build-dependencies]
 build-script-cfg.workspace = true
 search-cuda-tools.workspace = true
-cc = "1.0"

devices/nvidia-gpu/Cargo.toml

@@ -1,17 +1,13 @@
-fn main() {
+fn main() {
     use build_script_cfg::Cfg;
-    use search_cuda_tools::find_cuda_root;
+    use search_cuda_tools::{find_cuda_root, find_nccl_root};
 
     let cuda = Cfg::new("detected_cuda");
+    let nccl = Cfg::new("detected_nccl");
     if find_cuda_root().is_some() {
         cuda.define();
-        println!("cargo:rerun-if-changed=src/sample.cu");
-        cc::Build::new()
-            .cuda(true)
-            .flag("-gencode")
-            .flag("arch=compute_80,code=sm_80")
-            .flag("-allow-unsupported-compiler")
-            .file("src/sample.cu")
-            .compile("sample");
+        if find_nccl_root().is_some() {
+            nccl.define();
+        }
     }
 }