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

@@ -39,7 +39,7 @@ fn test_infer() {
     } = meta;
 
     let time = Instant::now();
-    let image = Image::load(picture);
+    let image = Image::load(&picture);
     println!("load image {:?}", time.elapsed());
 
     let time = Instant::now();
@@ -48,38 +48,55 @@ fn test_infer() {
         .normalize(dt, image_mean, image_std);
     println!("slice image {:?}", time.elapsed());
 
+    let batch = slices.batch();
+    let mut img_embd = meta.projector.img_embd(meta.dt, batch).map(Blob::new);
+    let d = img_embd.shape()[2];
+
     let weights = Weights::new(&storage);
     let mut worker = Worker::new(&Cpu, meta.clone(), weights);
 
-    let whole = slices.whole();
-    worker
-        .launch(
-            ClipArgs {
-                raw: whole.to_nchw(),
-                pos: pos70(whole.shape(), d_patch).map_slice(),
-                pos_resampler: pos_resampler(3584, whole.shape(), d_patch).map_slice(),
-            },
-            &mut [],
-            &ThisThread,
-        )
-        .unwrap();
+    {
+        let whole = slices.whole();
+        let img_embd = img_embd.map_slice_mut().slice(0, 0, 1, 1);
+        worker
+            .launch(
+                ClipArgs {
+                    img_embd,
+                    raw: whole.to_nchw(),
+                    pos: pos70(whole.shape(), d_patch).map_slice(),
+                    pos_resampler: pos_resampler(d, whole.shape(), d_patch).map_slice(),
+                },
+                &mut [],
+                &ThisThread,
+            )
+            .unwrap();
+    }
 
     if let Some(patches) = slices.patches_nchw() {
         let &[_, 3, h, w] = patches.shape() else {
             unreachable!()
         };
+        let img_embd = img_embd.map_slice_mut().slice(0, 1, 1, batch - 1);
         worker
             .launch(
                 ClipArgs {
+                    img_embd,
                     raw: patches.map_slice(),
                     pos: pos70([w, h], d_patch).map_slice(),
-                    pos_resampler: pos_resampler(3584, [w, h], d_patch).map_slice(),
+                    pos_resampler: pos_resampler(d, [w, h], d_patch).map_slice(),
                 },
                 &mut [],
                 &ThisThread,
             )
             .unwrap();
     }
+
+    println!(
+        "create {} x {} tokens from {}",
+        img_embd.shape()[0],
+        img_embd.shape()[1],
+        picture.display(),
+    );
 }
 
 fn pos70([w, h]: [usize; 2], d_patch: usize) -> Tensor<Blob> {

models/clip/common/src/args.rs

@@ -2,6 +2,8 @@
 use tensor::Tensor;
 
 pub struct Args<'a, H: Hardware> {
+    /// shape: [batch, 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.merge(0..2).unwrap();
                 {
                     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/image.rs

@@ -184,6 +184,12 @@ impl ImageGrid {
         }
     }
 
+    #[inline]
+    pub fn batch(&self) -> usize {
+        let [x, y] = self.grid();
+        x * y + 1
+    }
+
     pub fn patch(&self, x: usize, y: usize) -> Image<&[u8]> {
         Image(
             self.grid

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 {