feat(llama.cu): 根据维度选 add 算子

Author: CearX

Cargo.lock

@@ -5,7 +5,7 @@ edition.workspace = true
 
 [dependencies]
 operators = { git = "https://github.com/CearX/operators-rs.git", rev = "8a0d58a", default-features = false, features = ["nvidia-gpu", "common-cpu"] }
-nn = { git = "https://github.com/CearX/InfiniNN.git", rev = "74b1bea" }
+nn = { git = "https://github.com/CearX/InfiniNN.git", rev = "3ba7418" }
 ggus = { git = "https://github.com/InfiniTensor/gguf", rev = "23c362f" }
 tokeneer = { git = "https://github.com/InfiniTensor/tokeneer", rev = "c48f39f" }
 

llama.cu/Cargo.toml

@@ -453,7 +453,7 @@ impl<'ctx> Internal<'ctx> {
 fn builder() -> GraphBuilder {
     let mut ans = GraphBuilder::default();
     ans.register_op("embedding", op::embedding::Embedding)
-        .register_op("add4d", op::add4d::Add4d)
+        .register_op("add", op::add::Add)
         .register_op("conv", op::conv::Conv)
         .register_op("layer-norm", op::normalization::LayerNorm)
         .register_op("rms-norm", op::normalization::RmsNorm)

llama.cu/src/exec/group.rs

@@ -196,7 +196,11 @@ impl<'ctx> Handle<'ctx> {
             "rms-norm" => launch!(RmsNorm),
             "layer-norm" => launch!(LayerNorm),
             "linear" => launch!(Linear),
-            "add4d" => launch!(Add4d),
+            "add" => match inputs[0].shape().len() {
+                2 => launch!(Add),
+                4 => launch!(Add4d),
+                _ => panic!("add: unsupported shape"),
+            },
             "rope" => launch!(Rope),
             "mrope" => launch!(MRope),
             "gelu" => launch!(Gelu),

llama.cu/src/exec/step.rs

@@ -21,6 +21,7 @@ use operators::cuda::{Stream, VirByte};
 
 pub mod random_sample;
 
+pub use add::Add;
 pub use add4d::Add4d;
 #[cfg(nccl)]
 pub use all_reduce::AllReduce;