fix(xtask): nvidia 启用时退出需要同步以免释放异常

Signed-off-by: YdrMaster <[email protected]>
InfiniTensor · Apr 25, 2024 · 00778a2 · 00778a2
1 parent 0abf47a
commit 00778a2
Show file tree

Hide file tree

Showing 4 changed files with 10 additions and 3 deletions.
diff --git a/nvidia/common/src/rms_norm.rs b/nvidia/common/src/rms_norm.rs
@@ -61,7 +61,7 @@ extern "C" __global__ void {folding}(
 
         let (ptx, log) = Ptx::compile(code);
         if !log.is_empty() {
-            println!("{log}");
+            warn!("{log}");
         }
         Self {
             ptx: ptx.unwrap(),

diff --git a/nvidia/transformer/src/lib.rs b/nvidia/transformer/src/lib.rs
@@ -20,7 +20,7 @@ use std::{
 };
 use transformer::{Kernels, Llama2, Memory};
 
-pub use common_nv::cuda;
+pub use common_nv::{cuda, synchronize};
 
 pub struct Transformer {
     host: Memory,

diff --git a/xtask/src/chat.rs b/xtask/src/chat.rs
@@ -38,7 +38,6 @@ impl InferenceArgs {
             #[cfg(feature = "nvidia")]
             &[n] => {
                 use transformer_nv::{cuda, Transformer as M};
-                cuda::init();
                 chat!(M; cuda::Device::new(n));
             }
             #[cfg(feature = "nvidia")]

diff --git a/xtask/src/main.rs b/xtask/src/main.rs
@@ -19,9 +19,17 @@ fn main() {
 
 #[inline]
 fn block_on(f: impl Future) {
+    #[cfg(feature = "nvidia")]
+    {
+        transformer_nv::cuda::init();
+    }
     let runtime = tokio::runtime::Runtime::new().unwrap();
     runtime.block_on(f);
     runtime.shutdown_background();
+    #[cfg(feature = "nvidia")]
+    {
+        transformer_nv::synchronize();
+    }
 }
 
 #[derive(Parser)]