feat(llama-nv): 简化 NV 推理

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

Author: YdrMaster

Cargo.toml

@@ -28,7 +28,7 @@ itertools = "0.13"
 build-script-cfg = "0.0"
 
 ndarray-layout = { git = "https://github.com/YdrMaster/ndarray-layout", rev = "f1fdd24" }
-operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "8c2227a", default-features = false }
+operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "923949f", default-features = false }
 
 search-cl-tools = { git = "https://github.com/InfiniTensor/clrt", rev = "6846d52" }
 search-infini-tools = { git = "https://github.com/InfiniTensor/infini-rt", rev = "136c30b" }

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

@@ -8,7 +8,7 @@ use operators::{
     nvidia_gpu::{Config, Gpu},
     random_sample::{KVPair, SampleArgs},
 };
-use std::{slice::from_raw_parts_mut, thread, usize};
+use std::{slice::from_raw_parts_mut, time::Instant, usize};
 use test_utils::{load_roll_cache_size, Inference, TokenizerAndPrompt};
 
 type Worker<'w> = LlamaWorker<Operators, Weights<'w>>;
@@ -60,73 +60,66 @@ fn test_infer() {
         ..
     } = meta;
 
-    thread::scope(|s| {
-        let sample = s.spawn(move || {
-            let mut sample = RandomSample::new(gpu);
-            sample.scheme(dt_embd, nvoc).unwrap();
-            sample
-        });
-        gpu.apply(|ctx| {
-            let stream = ctx.stream();
-
-            let token_embd = stream.from_host(model.token_embd);
-            let weights = Weights::new(&model, .., 1, roll_cache_size, ctx);
-            let mut worker = Worker::new(&gpu, meta.clone(), weights, true);
-            let mut cache = meta.kv_cache(nctx).map(|size| stream.malloc::<u8>(size));
-            let sin_cos =
-                <Operators as llama::Operators>::build_sin_cos(dt_embd, nctx, dh, &stream);
-            let indices = RandomSample::build_indices(nvoc, &stream);
-
-            let sample = sample.join().unwrap();
-            test_utils::test_infer(eos, tokenizer, &prompt, max_steps, |input, pos| {
-                let mut embd = meta.embd(input.len()).map(|len| stream.malloc::<u8>(len));
-                let mut logits = meta.logits(1).map(|len| stream.malloc::<u8>(len));
-
-                let d = embd.get().len() / input.len();
-                for (i, &tok) in input.iter().enumerate() {
-                    stream.memcpy_d2d(
-                        &mut embd.get_mut()[i * d..][..d],
-                        &token_embd[tok as usize * d..][..d],
-                    )
-                }
-
-                worker
-                    .launch(
-                        LlamaArgs {
-                            embd: embd.map_slice_mut(),
-                            logits: logits.map_slice_mut(),
-                            sin_cos: sin_cos.map_slice(),
-                            requests: vec![LlamaRequest {
-                                cache: cache.map_slice_mut(),
-                                seq_len: input.len(),
-                                out_len: 1,
-                                pos,
-                            }],
-                            num_tokens: input.len(),
-                            max_seq_len: input.len(),
-                            max_att_len: pos + input.len(),
-                        },
-                        &mut [],
-                        &stream,
-                    )
-                    .unwrap();
-
-                let mut pairs = Tensor::kv_pair_vec(1, |size| stream.malloc::<u8>(size));
-
-                sample
-                    .launch(&mut pairs, &logits, &indices, sample_args, &mut [], &stream)
-                    .unwrap();
-
-                let mut pair = KVPair::new(0, f16::ZERO);
-                memcpy_d2h(
-                    unsafe {
-                        from_raw_parts_mut(&mut pair as *mut _ as *mut u8, size_of_val(&pair))
+    gpu.apply(|ctx| {
+        let stream = ctx.stream();
+
+        let time = Instant::now();
+        let token_embd = stream.from_host(model.token_embd);
+        let weights = Weights::new(&model, .., 1, roll_cache_size, ctx);
+        println!("load weights: {:?}", time.elapsed());
+
+        let mut worker = Worker::new(&gpu, meta.clone(), weights, true);
+        let mut cache = meta.kv_cache(nctx).map(|size| stream.malloc::<u8>(size));
+        let sin_cos = <Operators as llama::Operators>::build_sin_cos(dt_embd, nctx, dh, &stream);
+        let indices = RandomSample::build_indices(nvoc, &stream);
+        let sample = RandomSample::new(gpu);
+
+        test_utils::test_infer(eos, tokenizer, &prompt, max_steps, |input, pos| {
+            let mut embd = meta.embd(input.len()).map(|len| stream.malloc::<u8>(len));
+            let mut logits = meta.logits(1).map(|len| stream.malloc::<u8>(len));
+
+            let d = embd.get().len() / input.len();
+            for (i, &tok) in input.iter().enumerate() {
+                stream.memcpy_d2d(
+                    &mut embd.get_mut()[i * d..][..d],
+                    &token_embd[tok as usize * d..][..d],
+                )
+            }
+
+            worker
+                .launch(
+                    LlamaArgs {
+                        embd: embd.map_slice_mut(),
+                        logits: logits.map_slice_mut(),
+                        sin_cos: sin_cos.map_slice(),
+                        requests: vec![LlamaRequest {
+                            cache: cache.map_slice_mut(),
+                            seq_len: input.len(),
+                            out_len: 1,
+                            pos,
+                        }],
+                        num_tokens: input.len(),
+                        max_seq_len: input.len(),
+                        max_att_len: pos + input.len(),
                     },
-                    pairs.get(),
-                );
+                    &mut [],
+                    &stream,
+                )
+                .unwrap();
+
+            let mut pairs = Tensor::kv_pair_vec(1, |size| stream.malloc::<u8>(size));
+
+            sample
+                .launch(&mut pairs, &logits, &indices, sample_args, &mut [], &stream)
+                .unwrap();
+
+            let mut pair = KVPair::new(0, f16::ZERO);
+            memcpy_d2h(
+                unsafe { from_raw_parts_mut(&mut pair as *mut _ as *mut u8, size_of_val(&pair)) },
+                pairs.get(),
+            );
 
-                pair.idx() as _
-            });
+            pair.idx() as _
         });
     });
 }