perf(nv): 加速大模型加载

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

Author: YdrMaster

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

@@ -7,7 +7,7 @@ use operators::{
     nvidia_gpu::Gpu,
     random_sample::nvidia_gpu::Operator as RandomSampleGpu,
     rearrange::nvidia_gpu::Operator as Rearrange,
-    ByteOf, QueueOf, TopoNode,
+    Blob, ByteOf, QueueOf, TopoNode,
 };
 use std::{
     cell::{RefCell, RefMut},
@@ -207,9 +207,7 @@ impl<'blk> Weights<'blk> {
                 .as_ref()
                 .map(|_| Vec::with_capacity(model.meta.nblk));
             for blk in &model.blocks {
-                let blk = blk.distribute(&model.meta, range.clone(), count, |len| {
-                    ctx.malloc_host::<u8>(len)
-                });
+                let blk = blk.distribute(&model.meta, range.clone(), count, Blob::new);
                 let loader = loader
                     .get_or_insert_with(|| blk.as_ref().map(|s| H2DLoader::new(s.len(), &stream)));
 
@@ -240,20 +238,20 @@ impl<'blk> Weights<'blk> {
 
 struct H2DLoader<'ctx> {
     event: Event<'ctx>,
-    host: HostMem<'ctx>,
+    host: Blob,
     dev: DevMem<'ctx>,
 }
 
 impl<'ctx> H2DLoader<'ctx> {
     fn new(size: usize, stream: &Stream<'ctx>) -> Self {
         Self {
             event: stream.record(),
-            host: stream.ctx().malloc_host::<u8>(size),
+            host: Blob::new(size),
             dev: stream.malloc::<u8>(size),
         }
     }
 
-    fn load(&mut self, host: Contiguous<HostMem<'ctx>>, stream: &Stream<'ctx>) -> DevMem<'ctx> {
+    fn load(&mut self, host: Contiguous<Blob>, stream: &Stream<'ctx>) -> DevMem<'ctx> {
         self.event.synchronize();
         match host {
             Contiguous::Borrowed(host) => self.host.copy_from_slice(host),

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

@@ -14,7 +14,7 @@ use regex::Regex;
 use std::{
     iter::zip,
     slice::{from_raw_parts, from_raw_parts_mut},
-    thread,
+    thread, u64,
 };
 use test_utils::{test_infer_paralle, Inference, Task, TokenizerAndPrompt, WorkerSeed};
 
@@ -88,6 +88,14 @@ fn test_infer() {
                     let WorkerSeed { node, tasks } = seed;
                     node.processor().apply(|ctx| {
                         let stream = ctx.stream();
+
+                        let mut free = 0;
+                        let mut total = 0;
+                        cuda::driver!(cuMemGetInfo_v2(&mut free, &mut total));
+
+                        ctx.dev().set_mempool_threshold(u64::MAX);
+                        let _ = stream.malloc::<u8>((free >> 30).saturating_sub(1) << 30);
+
                         info!("worker[{id}] loading weights...");
                         let weights = Weights::new(model, range, count, usize::MAX, ctx);
                         let mut worker = Worker::new(id, &node, meta.clone(), weights, id == 0);