feat(clip): 完善 clip 结构并为 softmax 添加并行加速

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

Author: YdrMaster

Cargo.toml

@@ -38,7 +38,7 @@ itertools = "0.13"
 env_logger = "0.11"
 build-script-cfg = "0.0"
 
-operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "0bd4107", default-features = false }
+operators = { git = "https://github.com/YdrMaster/operators-rs", rev = "8712870", default-features = false }
 
 search-cl-tools = { git = "https://github.com/InfiniTensor/clrt", rev = "f69b160" }
 search-infini-tools = { git = "https://github.com/InfiniTensor/infini-rt", rev = "e8362c3" }

models/clip/common-cpu/src/infer.rs

@@ -52,9 +52,11 @@ fn test_infer() {
     let mut worker = Worker::new(&Cpu, meta.clone(), weights);
 
     let whole = slices.whole();
+    let mut img_embd = meta.projector.img_embd(meta.dt, 1).map(Blob::new);
     worker
         .launch(
             ClipArgs {
+                img_embd: img_embd.map_slice_mut(),
                 raw: whole.to_nchw(),
                 pos: pos70(whole.shape(), d_patch).map_slice(),
                 pos_resampler: pos_resampler(3584, whole.shape(), d_patch).map_slice(),
@@ -65,12 +67,14 @@ fn test_infer() {
         .unwrap();
 
     if let Some(patches) = slices.patches_nchw() {
-        let &[_, 3, h, w] = patches.shape() else {
+        let &[batch, 3, h, w] = patches.shape() else {
             unreachable!()
         };
+        let mut img_embd = meta.projector.img_embd(meta.dt, batch).map(Blob::new);
         worker
             .launch(
                 ClipArgs {
+                    img_embd: img_embd.map_slice_mut(),
                     raw: patches.map_slice(),
                     pos: pos70([w, h], d_patch).map_slice(),
                     pos_resampler: pos_resampler(3584, [w, h], d_patch).map_slice(),

models/clip/common/src/args.rs

@@ -2,6 +2,8 @@
 use tensor::Tensor;
 
 pub struct Args<'a, H: Hardware> {
+    /// shape: [batch x projector_dp, projector_d]
+    pub img_embd: Tensor<&'a mut [H::Byte]>,
     /// shape: [n, c, h, w]
     pub raw: Tensor<&'a [H::Byte]>,
     /// shape: [h x w]

models/clip/common/src/compute.rs

@@ -149,6 +149,7 @@ where
     {
         let time = Instant::now();
         let Args {
+            img_embd: proj_q,
             raw,
             pos,
             pos_resampler,
@@ -317,10 +318,7 @@ where
                 let [w, b] = weights.resampler_attn_o(queue);
                 let attn_o = (attn_w.clone().map(|_| w), Some(attn_b.clone().map(|_| b)));
 
-                let qo = Tensor::new(dt, &[batch * dq, d]);
-
-                let (buf, workspace) = workspace.split_at_mut(*qo.get());
-                let mut q_ = qo.clone().map(|_| buf);
+                let mut q_ = proj_q;
                 {
                     let mut q_ = q_.map_slice_mut().tile(0, &[batch, dq]);
                     {
@@ -363,18 +361,19 @@ where
                 }
                 let o = q_;
 
-                let (buf, workspace) = workspace.split_at_mut(*qo.get());
-                let mut o_ = qo.map(|_| buf);
+                let o_ = Tensor::new(o.dt(), o.shape());
+                let (buf, workspace) = workspace.split_at_mut(*o_.get());
+                let mut o_ = o_.map(|_| buf);
                 self.mat_mul(&mut o_, &o, attn_o, workspace, queue_alloc)?;
 
                 let [w, b] = weights.resampler_ln_post(queue);
                 let ln_post = [ln.clone().map(|_| w), ln.clone().map(|_| b)];
                 let inplace = unsafe { o_.map_slice_static() };
                 self.layer_norm(&mut o_, &inplace, ln_post, workspace, queue_alloc)?;
 
-                let mut out = o;
+                let mut img_embd = o;
                 let w = attn_w.map(|_| weights.resampler_proj(queue));
-                self.mat_mul(&mut out, &o_, (w, None), workspace, queue_alloc)?
+                self.mat_mul(&mut img_embd, &o_, (w, None), workspace, queue_alloc)?
             }
         }
 

models/clip/common/src/projector/mod.rs

@@ -1,6 +1,7 @@
 pub(crate) mod resampler;
 
-use gguf::{GGufMetaMapExt, GGufModel};
+use gguf::{ggml_quants::digit_layout::DigitLayout, GGufMetaMapExt, GGufModel};
+use tensor::Tensor;
 
 #[derive(Clone, Debug)]
 pub enum ProjectorMeta {
@@ -14,6 +15,12 @@ impl ProjectorMeta {
             projector => todo!("unsupported projector type: {projector}"),
         }
     }
+
+    pub fn img_embd(&self, dt: DigitLayout, batch: usize) -> Tensor<usize> {
+        match self {
+            ProjectorMeta::Resampler(meta) => meta.img_embd(dt, batch),
+        }
+    }
 }
 
 #[derive(Clone)]

models/clip/common/src/projector/resampler.rs

@@ -1,4 +1,5 @@
-use gguf::{tensor, GGufMetaMapExt, GGufModel};
+use gguf::{ggml_quants::digit_layout::DigitLayout, tensor, GGufMetaMapExt, GGufModel};
+use tensor::Tensor;
 
 #[derive(Clone, Debug)]
 pub struct Meta {
@@ -23,6 +24,11 @@ impl Meta {
             version => todo!("Unsupported MiniCPM version: {version}"),
         }
     }
+
+    #[inline]
+    pub fn img_embd(&self, dt: DigitLayout, batch: usize) -> Tensor<usize> {
+        Tensor::new(dt, &[batch * self.dq, self.d])
+    }
 }
 
 #[derive(Clone)]

tensor/src/lib.rs

@@ -84,7 +84,7 @@ impl<T> Tensor<T> {
 
         let merged = self
             .layout
-            .merge_be(0, self.layout.ndim())
+            .merge_free(0, self.layout.ndim())
             .expect("dense tensor is castable");
         let &[d] = merged.shape() else { unreachable!() };
         let &[s] = merged.strides() else {