perf: 优化性能

Signed-off-by: YdrMaster <[email protected]>

Author: YdrMaster

Cargo.lock

@@ -4,10 +4,10 @@ version = "0.0.0"
 edition.workspace = true
 
 [dependencies]
-cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "803d64b" }
-cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "803d64b" }
-nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "803d64b" }
-flash-attn = { git = "https://github.com/YdrMaster/learn-flash-attn", rev = "1caeeee" }
+cuda = { git = "https://github.com/YdrMaster/cuda-driver", rev = "6a97931" }
+cublas = { git = "https://github.com/YdrMaster/cuda-driver", rev = "6a97931" }
+nccl = { git = "https://github.com/YdrMaster/cuda-driver", rev = "6a97931" }
+flash-attn = { git = "https://github.com/YdrMaster/learn-flash-attn", rev = "616bbac" }
 nn = { git = "https://github.com/YdrMaster/InfiniNN", rev = "171c5b0" }
 ggus = { git = "https://github.com/InfiniTensor/gguf", rev = "23c362f" }
 tokeneer = { git = "https://github.com/InfiniTensor/tokeneer", rev = "c48f39f" }

llama.cu/Cargo.toml

@@ -147,11 +147,19 @@ impl<'ctx> Handle<'ctx> {
         use ::flash_attn::attention::cuda::code as flash_attn_code;
         let Attention { q, k, v, o, .. } = attn;
         let dt = distinct(&[q.dt(), k.dt(), v.dt(), o.dt()]).unwrap();
+        let d = *q.shape().last().unwrap();
         // 编译
-        let key = [ModuleKey::Text("flash-attn"), ModuleKey::Type(dt)].into_iter();
+        let key = [
+            ModuleKey::Text("flash-attn"),
+            ModuleKey::Type(dt),
+            ModuleKey::Size(d),
+        ]
+        .into_iter();
         match dt {
             types::F16 => {
-                let module = self.compile(key.collect(), || flash_attn_code::<f16>());
+                let module = self.compile(key.collect(), || {
+                    flash_attn_code::<f16>(d, stream.ctx().dev().warp_size())
+                });
                 launch_attn_typed::<f16>(attn, reqs, module, stream)
             }
             _ => todo!(),

llama.cu/src/exec/step.rs

@@ -1,6 +1,6 @@
 use crate::op::ModuleKey;
 use cublas::Cublas;
-use cuda::{CurrentCtx, Module, Ptx};
+use cuda::{CurrentCtx, Module, Rtc};
 use std::collections::HashMap;
 
 #[cfg(nccl)]
@@ -37,9 +37,11 @@ impl<'ctx> Handle<'ctx> {
 
     pub fn compile(&mut self, key: Box<[ModuleKey]>, code: impl FnOnce() -> String) -> &Module {
         self.modules.entry(key).or_insert_with(|| {
-            let (ptx, log) = Ptx::compile(code(), self.ctx.dev().compute_capability());
-            let Ok(ptx) = ptx else { panic!("{log}") };
-            self.ctx.load(&ptx)
+            let program = Rtc::new()
+                .arch(self.ctx.dev().compute_capability())
+                .compile(&code())
+                .unwrap();
+            self.ctx.load(&program)
         })
     }
 

llama.cu/src/handle.rs

@@ -1,8 +1,7 @@
 //! <https://zhuanlan.zhihu.com/p/667025336>
 
 use crate::utils::offset_ptr;
-use cuda::{CurrentCtx, DevByte, DevMem, Module, Ptx, Stream, VirByte, params};
-use log::warn;
+use cuda::{CurrentCtx, DevByte, DevMem, Module, Rtc, Stream, VirByte, params};
 use nn::Tensor;
 use std::ffi::c_uint;
 use tokeneer::utok;
@@ -77,15 +76,10 @@ extern "C" __global__ void next(
     next_kernel(logits, records, n, eos, temperature, penalty, tok);
 }}"#
         );
-        let (ptx, log) = Ptx::compile(code, ctx.dev().compute_capability());
-        match ptx {
-            Ok(ptx) => {
-                if !log.is_empty() {
-                    warn!("{log}")
-                }
-                ctx.load(&ptx)
-            }
-            Err(e) => panic!("logits modify compilation failed with {e:?}, log:\n {log}"),
-        }
+        let program = Rtc::new()
+            .arch(ctx.dev().compute_capability())
+            .compile(&code)
+            .unwrap();
+        ctx.load(&program)
     }
 }

llama.cu/src/op/random_sample/modifier.rs

@@ -35,7 +35,7 @@ fn main() {
 
 #[derive(Parser)]
 #[clap(name = "InfiniLM")]
-#[clap(version, about, long_about = None)]
+#[clap(version, long_about = None)]
 struct Cli {
     #[clap(subcommand)]
     command: Commands,