add minicpm with a bug

src/lib.rs

@@ -1,4 +1,4 @@
-use crate::models::{GenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel};
+use crate::models::{minicpm4::generate::MiniCPMGenerateModel, qwen2_5vl::generate::Qwen2_5VLGenerateModel, GenerateModel};
 use anyhow::{Ok, Result};
 use candle_core::{DType, Device};
 pub mod chat_template;
@@ -9,6 +9,7 @@ pub mod utils;
 
 pub enum ModelType {
     Qwen2_5VL,
+    MiniCPM4,
 }
 
 impl ModelType {
@@ -22,6 +23,10 @@ impl ModelType {
             ModelType::Qwen2_5VL => {
                 let model = Qwen2_5VLGenerateModel::init(model_path, device, dtype)?;
                 Ok(Box::new(model))
+            },
+            ModelType::MiniCPM4 => {
+                let model = MiniCPMGenerateModel::init(model_path, device, dtype)?;
+                Ok(Box::new(model))  
             }
         }
     }

src/models/base_modules/mod.rs

@@ -0,0 +1,262 @@
+use anyhow::Result;
+use candle_core::{Tensor, D};
+use candle_nn::{Activation, Linear, Module, VarBuilder, linear, linear_no_bias};
+
+use crate::{position_embed::rope::apply_rotary_pos_emb, utils::tensor_utils::repeat_kv};
+
+#[derive(Debug, Clone)]
+pub struct MLPWithBias {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl MLPWithBias {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        intermediate_size: usize,
+        act_fn: Activation,
+    ) -> Result<Self> {
+        let gate_proj = linear(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
+        let up_proj = linear(hidden_size, intermediate_size, vb.pp("up_proj"))?;
+        let down_proj = linear(intermediate_size, hidden_size, vb.pp("down_proj"))?;
+        Ok(Self {
+            gate_proj,
+            up_proj,
+            down_proj,
+            act_fn,
+        })
+    }
+}
+
+impl Module for MLPWithBias {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?;
+        let rhs = xs.apply(&self.up_proj)?;
+        (lhs * rhs)?.apply(&self.down_proj)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct MLPNoBias {
+    gate_proj: Linear,
+    up_proj: Linear,
+    down_proj: Linear,
+    act_fn: Activation,
+}
+
+impl MLPNoBias {
+    pub fn new(
+        vb: VarBuilder,
+        hidden_size: usize,
+        intermediate_size: usize,
+        act_fn: Activation,
+    ) -> Result<Self> {
+        let gate_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("gate_proj"))?;
+        let up_proj = linear_no_bias(hidden_size, intermediate_size, vb.pp("up_proj"))?;
+        let down_proj = linear_no_bias(intermediate_size, hidden_size, vb.pp("down_proj"))?;
+        Ok(Self {
+            gate_proj,
+            up_proj,
+            down_proj,
+            act_fn,
+        })
+    }
+}
+
+impl Module for MLPNoBias {
+    fn forward(&self, xs: &Tensor) -> candle_core::Result<Tensor> {
+        let lhs = xs.apply(&self.gate_proj)?.apply(&self.act_fn)?;
+        let rhs = xs.apply(&self.up_proj)?;
+        (lhs * rhs)?.apply(&self.down_proj)
+    }
+}
+
+#[derive(Debug, Clone)]
+pub struct AttentionNobias {
+    q_proj: Linear,
+    k_proj: Linear,
+    v_proj: Linear,
+    o_proj: Linear,
+    num_heads: usize,
+    num_kv_heads: usize,
+    num_kv_groups: usize,
+    head_dim: usize,
+    hidden_size: usize,
+    kv_cache: Option<(Tensor, Tensor)>,
+}
+
+impl AttentionNobias {
+    pub fn new(vb: VarBuilder, hidden_size: usize, num_attention_heads: usize, num_key_value_heads: usize) -> Result<Self> {
+        let num_kv_groups = num_attention_heads / num_key_value_heads;
+        let head_dim = hidden_size / num_attention_heads;
+        let q_proj = linear_no_bias(hidden_size, num_attention_heads * head_dim, vb.pp("q_proj"))?;
+        let k_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("k_proj"))?;
+        let v_proj = linear_no_bias(hidden_size, num_key_value_heads * head_dim, vb.pp("v_proj"))?;
+        let o_proj = linear_no_bias(hidden_size, hidden_size, vb.pp("o_proj"))?;
+        Ok(Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            num_heads: num_attention_heads,
+            num_kv_heads: num_key_value_heads,
+            num_kv_groups,
+            head_dim,
+            hidden_size,
+            kv_cache: None,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (b_sz, q_len, _) = xs.dims3()?;
+        let query_states = self.q_proj.forward(xs)?;
+        let key_states = self.k_proj.forward(xs)?;
+        let value_states = self.v_proj.forward(xs)?;
+        let query_states = query_states
+            .reshape((b_sz, q_len, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let key_states = key_states
+            .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let value_states = value_states
+            .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let (query_states, key_states) =
+            apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?;
+
+        let key_states = repeat_kv(key_states, self.num_kv_groups)?.contiguous()?;
+        let value_states = repeat_kv(value_states, self.num_kv_groups)?.contiguous()?;
+        let query_states = query_states.contiguous()?;
+        let attn_output = {
+            let scale = 1f64 / f64::sqrt(self.head_dim as f64);
+            #[cfg(not(feature = "flash-attn"))]
+            {
+                let attn_weights =
+                    query_states.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?;
+                let attn_weights = (attn_weights * scale)?;
+                let attn_weights = match attention_mask {
+                    None => attn_weights,
+                    Some(mask) => attn_weights.broadcast_add(mask)?,
+                };
+                let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
+                let attn_weights = attn_weights.matmul(&value_states)?;
+                attn_weights
+            }
+            #[cfg(feature = "flash-attn")]
+            {
+                // use flash-attn,
+                // flash-attn shape: (bs, seq_len, num_head, head_dim)
+                let query_states = query_states.transpose(1, 2)?;
+                let key_states = key_states.transpose(1, 2)?;
+                let value_states = value_states.transpose(1, 2)?;
+                let attn_output = candle_flash_attn::flash_attn(
+                    &query_states,
+                    &key_states,
+                    &value_states,
+                    scale as f32,
+                    attention_mask.is_some(),
+                )?
+                .transpose(1, 2)?;
+                attn_output
+            }
+        };
+        let attn_output =
+            attn_output
+                .transpose(1, 2)?
+                .contiguous()?
+                .reshape((b_sz, q_len, self.hidden_size))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn forward_step(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let (b_sz, q_len, _) = xs.dims3()?;
+        let query_states = self.q_proj.forward(xs)?;
+        let key_states = self.k_proj.forward(xs)?;
+        let value_states = self.v_proj.forward(xs)?;
+        let query_states = query_states
+            .reshape((b_sz, q_len, self.num_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let key_states = key_states
+            .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let value_states = value_states
+            .reshape((b_sz, q_len, self.num_kv_heads, self.head_dim))?
+            .transpose(1, 2)?;
+        let (query_states, key_states) =
+            apply_rotary_pos_emb(&query_states, &key_states, cos, sin)?;
+        let (key_states, value_states) = match &self.kv_cache {
+            None => (key_states, value_states),
+            Some((prev_k, prev_v)) => {
+                let key_states = Tensor::cat(&[prev_k, &key_states], 2)?;
+                let value_states = Tensor::cat(&[prev_v, &value_states], 2)?;
+                (key_states, value_states)
+            }
+        };
+
+        self.kv_cache = Some((key_states.clone(), value_states.clone()));
+
+        let key_states = repeat_kv(key_states, self.num_kv_groups)?.contiguous()?;
+        let value_states = repeat_kv(value_states, self.num_kv_groups)?.contiguous()?;
+        let query_states = query_states.contiguous()?;
+        let attn_output = {
+            let scale = 1f64 / f64::sqrt(self.head_dim as f64);
+            #[cfg(not(feature = "flash-attn"))]
+            {
+                let attn_weights =
+                    query_states.matmul(&key_states.transpose(D::Minus2, D::Minus1)?)?;
+                let attn_weights = (attn_weights * scale)?;
+                let attn_weights = match attention_mask {
+                    None => attn_weights,
+                    Some(mask) => attn_weights.broadcast_add(mask)?,
+                };
+                let attn_weights = candle_nn::ops::softmax_last_dim(&attn_weights)?;
+                let attn_weights = attn_weights.matmul(&value_states)?;
+                attn_weights
+            }
+            #[cfg(feature = "flash-attn")]
+            {
+                // use flash-attn,
+                // flash-attn shape: (bs, seq_len, num_head, head_dim)
+                let query_states = query_states.transpose(1, 2)?;
+                let key_states = key_states.transpose(1, 2)?;
+                let value_states = value_states.transpose(1, 2)?;
+                let attn_output = candle_flash_attn::flash_attn(
+                    &query_states,
+                    &key_states,
+                    &value_states,
+                    scale as f32,
+                    attention_mask.is_some(),
+                )?
+                .transpose(1, 2)?;
+                attn_output
+            }
+        };
+        let attn_output =
+            attn_output
+                .transpose(1, 2)?
+                .contiguous()?
+                .reshape((b_sz, q_len, self.hidden_size))?;
+        let attn_output = attn_output.apply(&self.o_proj)?;
+        Ok(attn_output)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        self.kv_cache = None
+    }
+}

src/models/minicpm4/config.rs

@@ -0,0 +1,32 @@
+use candle_nn::Activation;
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct RopeScalingConfig {
+    pub rope_type: String,
+    pub long_factor: Vec<f32>,
+    pub short_factor: Vec<f32>,
+    pub original_max_position_embeddings: usize,
+}
+
+#[derive(Debug, Clone, PartialEq, serde::Deserialize)]
+pub struct MiniCPM4Config {
+    pub bos_token_id: u32,
+    pub eos_token_id: Vec<u32>,
+    pub hidden_act: Activation,
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub max_position_embeddings: usize,
+    pub num_attention_heads: usize,
+    pub num_hidden_layers: usize,
+    pub num_key_value_heads: usize,
+    pub rms_norm_eps: f64,
+    pub rope_scaling: RopeScalingConfig,
+    pub torch_dtype: String,
+    pub vocab_size: usize,
+    // pub use_mup: bool,
+    pub scale_emb:f32,
+    pub dim_model_base: usize,
+    pub scale_depth: f32,
+    // pub rope_theta: f32,
+    // pub kv_channels: i32,
+}

src/models/minicpm4/generate.rs

@@ -0,0 +1,136 @@
+use crate::models::minicpm4::config::MiniCPM4Config;
+use crate::models::minicpm4::model::MiniCPMModel;
+// use crate::models::GenerateStream;
+use crate::utils::utils::{
+    build_completion_chunk_response, build_completion_response, find_safetensors_files, get_device,
+    get_dtype, get_logit_processor,
+};
+use crate::{
+    chat_template::chat_template::ChatTemplate, models::GenerateModel,
+    tokenizer::tokenizer::TokenizerModel,
+};
+use anyhow::{Result, anyhow};
+use candle_core::{D, DType, Device, IndexOp, Tensor};
+use candle_nn::VarBuilder;
+use openai_dive::v1::resources::chat::{
+    ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
+};
+use rocket::async_stream::stream;
+use rocket::futures::Stream;
+
+pub struct MiniCPMGenerateModel<'a> {
+    chat_template: ChatTemplate<'a>,
+    tokenizer: TokenizerModel,
+    minicpm: MiniCPMModel,
+    device: Device,
+    endoftext_id: u32,
+    im_end_id: u32,
+}
+
+impl<'a> GenerateModel for MiniCPMGenerateModel<'a> {
+    fn init(path: &str, device: Option<&Device>, dtype: Option<DType>) -> Result<Self> {
+        let chat_template = ChatTemplate::init(path)?;
+        let tokenizer = TokenizerModel::init(path)?;
+        let config_path = path.to_string() + "/config.json";
+        let cfg: MiniCPM4Config = serde_json::from_slice(&std::fs::read(config_path)?)?;
+        let device = &get_device(device);
+        let cfg_dtype = cfg.torch_dtype.as_str();
+        let dtype = get_dtype(dtype, cfg_dtype);
+        let endoftext_id = cfg.eos_token_id[0];
+        let im_end_id = cfg.eos_token_id[1];
+        let model_list = find_safetensors_files(&path)?;
+        let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
+        let minicpm = MiniCPMModel::new(vb, cfg)?;
+
+        Ok(MiniCPMGenerateModel {
+            chat_template,
+            tokenizer,
+            minicpm,
+            device: device.clone(),
+            endoftext_id,
+            im_end_id,
+        })
+    }
+
+    fn generate(&mut self, mes: ChatCompletionParameters) -> Result<ChatCompletionResponse> {
+        let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p);
+        let mes_render = self.chat_template.apply_chat_template(&mes)?;
+        let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?;
+        let mut seq_len = input_ids.dim(1)?;
+        let mut seqlen_offset = 0;
+        let mut generate = Vec::new();
+        let sample_len = match mes.max_tokens {
+            Some(max) => max,
+            None => 512,
+        };
+        for _ in 0..sample_len {
+            let logits = self.minicpm.forward_step(&input_ids, seqlen_offset)?;
+            let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+            let next_token = logit_processor.sample(&logits)?;
+            generate.push(next_token);
+            if next_token == self.endoftext_id || next_token == self.im_end_id {
+                break;
+            }
+            seqlen_offset += seq_len;
+            seq_len = 1;
+            input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+        }
+        let res = self.tokenizer.token_decode(generate)?;
+        self.minicpm.clear_kv_cache();
+        let response = build_completion_response(res, "minicpm");
+        Ok(response)
+    }
+    fn generate_stream(
+        &mut self,
+        mes: ChatCompletionParameters,
+    ) -> Result<impl Stream<Item = Result<ChatCompletionChunkResponse, anyhow::Error>>> {
+        let mut logit_processor = get_logit_processor(mes.temperature, mes.top_p);
+        let mes_render = self.chat_template.apply_chat_template(&mes)?;
+        let mut input_ids = self.tokenizer.text_encode(mes_render, &self.device)?;
+        let mut seq_len = input_ids.dim(1)?;
+        let mut seqlen_offset = 0;
+        let sample_len = match mes.max_tokens {
+            Some(max) => max,
+            None => 512,
+        };
+        let stream = stream! {
+            let mut error_tokens = Vec::new();
+            for _ in 0..sample_len {
+                let logits = self.minicpm.forward_step(
+                    &input_ids,
+                    seqlen_offset,
+                )?;
+                let logits = logits.squeeze(0)?.squeeze(0)?.to_dtype(DType::F32)?;
+                let next_token = logit_processor.sample(&logits)?;
+                let mut decode_ids = Vec::new();
+                if error_tokens.len() > 0 {
+                    decode_ids.extend_from_slice(&error_tokens);
+                }
+                decode_ids.push(next_token);
+                let decoded_token = self.tokenizer.token_decode(decode_ids).map_err(|e| anyhow!(format!("stream decode error{}", e)))?;
+                if decoded_token.contains("�") {
+                    error_tokens.push(next_token);
+                    if error_tokens.len() > 3 {
+                        error_tokens.clear();
+                    }
+                    seqlen_offset += seq_len;
+                    seq_len = 1;
+                    input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+                    continue;
+                }
+                error_tokens.clear();
+                let chunk = build_completion_chunk_response(decoded_token, "minicpm", None, None);
+                yield Ok(chunk);
+                if next_token == self.endoftext_id || next_token == self.im_end_id {
+                    break;
+                }
+                seqlen_offset += seq_len;
+                seq_len = 1;
+                input_ids = Tensor::from_vec(vec![next_token], (1, 1), &self.device)?;
+
+            }
+            self.minicpm.clear_kv_cache();
+        };
+        Ok(stream)
+    }
+}

src/models/minicpm4/mod.rs

@@ -0,0 +1,3 @@
+pub mod config;
+pub mod model;
+pub mod generate;

src/models/minicpm4/model.rs

@@ -0,0 +1,257 @@
+use crate::{
+    models::{
+        base_modules::{AttentionNobias, MLPNoBias},
+        minicpm4::config::MiniCPM4Config,
+    },
+    position_embed::rope::compute_default_rope_parameters,
+    utils::tensor_utils::prepare_causal_attention_mask,
+};
+use anyhow::{Ok, Result};
+use candle_core::{D, DType, Device, Tensor, Var};
+use candle_nn::{embedding, rms_norm, Embedding, Linear, Module, RmsNorm, VarBuilder};
+
+pub struct MiniCPMLongRoPE {
+    head_dim: usize,
+    rope_theta: f32,
+    max_position_embeddings: usize,
+    short_factor: Vec<f32>,
+    long_factor: Vec<f32>,
+    original_max_position_embeddings: usize,
+    inv_freq: Tensor,
+    cos_cached: Tensor,
+    sin_cached: Tensor,
+}
+impl MiniCPMLongRoPE {
+    pub fn new(cfg: &MiniCPM4Config, device: &Device) -> Result<Self> {
+        let head_dim = cfg.hidden_size / cfg.num_attention_heads;
+        let rope_theta = 10000.0;
+        let max_position_embeddings = cfg.max_position_embeddings;
+        let short_factor = cfg.rope_scaling.short_factor.clone();
+        let long_factor = cfg.rope_scaling.short_factor.clone();
+        let original_max_position_embeddings = cfg.rope_scaling.original_max_position_embeddings;
+        let scale = max_position_embeddings / original_max_position_embeddings;
+        let scaling_factor =
+            (1.0 + (scale as f64).ln() + (original_max_position_embeddings as f64).ln()).sqrt();
+        let inv_freq = compute_default_rope_parameters(head_dim, rope_theta);
+        let inv_freq = Tensor::from_slice(&inv_freq, (1, inv_freq.len()), device)?;
+        let t = Tensor::arange(0.0_f32, max_position_embeddings as f32, device)?
+            .reshape((max_position_embeddings, 1))?;
+        // short_factor.len() = 32
+        // head_dim = 1024 / 16 = 64, inv_freq.len() = 32
+        let ext_factors = Tensor::from_slice(&short_factor, (1, short_factor.len()), device)?;
+        let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?;
+        // (seq_len, 1) matmul (1, 32) -> (seq_len, 32) * (1, 32)-> (seq_len, 32)
+        let freqs = t.matmul(&ext_factors)?.broadcast_mul(&inv_freq)?;
+
+        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;
+        let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?;
+        let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?;
+        Ok(Self {
+            head_dim,
+            rope_theta,
+            max_position_embeddings,
+            short_factor,
+            long_factor,
+            original_max_position_embeddings,
+            inv_freq,
+            cos_cached,
+            sin_cached,
+        })
+    }
+    pub fn update_cos_sin_cache(&mut self, seqlen: usize, device: &Device) -> Result<()> {
+        let t = Tensor::arange(0.0_f32, seqlen as f32, device)?.reshape((seqlen, 1))?;
+        let mut ext_factors =
+            Tensor::from_slice(&self.short_factor, (1, self.short_factor.len()), device)?;
+        if seqlen > self.original_max_position_embeddings {
+            ext_factors =
+                Tensor::from_slice(&self.long_factor, (1, self.long_factor.len()), device)?;
+        }
+        let ext_factors = Tensor::ones_like(&ext_factors)?.div(&ext_factors)?;
+        let freqs = t.matmul(&ext_factors)?.broadcast_mul(&self.inv_freq)?;
+        let emb = Tensor::cat(&[&freqs, &freqs], D::Minus1)?;
+        let scale = seqlen / self.original_max_position_embeddings;
+        let scaling_factor =
+            (1.0 + (scale as f64).ln() + (self.original_max_position_embeddings as f64).ln())
+                .sqrt();
+        let cos_cached = emb.cos()?.affine(scaling_factor, 0.0)?;
+        let sin_cached = emb.sin()?.affine(scaling_factor, 0.0)?;
+        self.cos_cached = cos_cached;
+        self.sin_cached = sin_cached;
+        Ok(())
+    }
+    pub fn forward(&self, pos_offset: usize, seqlen: usize) -> Result<(Tensor, Tensor)> {
+        let cos = self.cos_cached.narrow(0, pos_offset, seqlen)?;
+        let sin = self.sin_cached.narrow(0, pos_offset, seqlen)?;
+        Ok((cos, sin))
+    }
+}
+
+pub struct MiniCPMDecoderLayer {
+    self_attn: AttentionNobias,
+    mlp: MLPNoBias,
+    input_layernorm: RmsNorm,
+    post_attention_layernorm: RmsNorm,
+    scale_depth: f32,
+    num_hidden_layers: usize,
+}
+
+impl MiniCPMDecoderLayer {
+    pub fn new(vb: VarBuilder, cfg: &MiniCPM4Config) -> Result<Self> {
+        let self_attn = AttentionNobias::new(
+            vb.pp("self_attn"),
+            cfg.hidden_size,
+            cfg.num_attention_heads,
+            cfg.num_key_value_heads,
+        )?;
+        let mlp = MLPNoBias::new(
+            vb.pp("mlp"),
+            cfg.hidden_size,
+            cfg.intermediate_size,
+            cfg.hidden_act,
+        )?;
+        let input_layernorm =
+            rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("input_layernorm"))?;
+        let post_attention_layernorm = rms_norm(
+            cfg.hidden_size,
+            cfg.rms_norm_eps,
+            vb.pp("post_attention_layernorm"),
+        )?;
+        Ok(Self {
+            self_attn,
+            mlp,
+            input_layernorm,
+            post_attention_layernorm,
+            scale_depth: cfg.scale_depth,
+            num_hidden_layers: cfg.num_hidden_layers,
+        })
+    }
+
+    pub fn forward(
+        &self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs;
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self.self_attn.forward(&xs, cos, sin, attention_mask)?;
+        let xs = (xs + residual)?;
+        let residual = &xs;
+        let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
+        let xs = (residual + xs)?;
+        Ok(xs)
+    }
+
+    pub fn forward_step(
+        &mut self,
+        xs: &Tensor,
+        cos: &Tensor,
+        sin: &Tensor,
+        attention_mask: Option<&Tensor>,
+    ) -> Result<Tensor> {
+        let residual = xs;
+        let xs = self.input_layernorm.forward(xs)?;
+        let xs = self.self_attn.forward_step(&xs, cos, sin, attention_mask)?;
+        let xs = (xs + residual)?;
+        let residual = &xs;
+        let xs = xs.apply(&self.post_attention_layernorm)?.apply(&self.mlp)?;
+        let xs = (residual + xs)?;
+        Ok(xs)
+    }
+    pub fn clear_kv_cache(&mut self) {
+        self.self_attn.clear_kv_cache();
+    }
+}
+
+pub struct MiniCPMModel {
+    cfg: MiniCPM4Config,
+    embed_tokens: Embedding,
+    layers: Vec<MiniCPMDecoderLayer>,
+    norm: RmsNorm,
+    rope_emb: MiniCPMLongRoPE,
+    lm_head: Linear,
+}
+
+impl MiniCPMModel {
+    pub fn new(vb: VarBuilder, cfg: MiniCPM4Config) -> Result<Self> {
+        let embed_tokens = embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?;
+        let mut layers = Vec::with_capacity(cfg.num_hidden_layers);
+        let vb_layers = vb.pp("layers");
+        for i in 0..cfg.num_hidden_layers {
+            let layer = MiniCPMDecoderLayer::new(vb_layers.pp(i), &cfg)?;
+            layers.push(layer);
+        }
+        let norm = rms_norm(cfg.hidden_size, cfg.rms_norm_eps, vb.pp("norm"))?;
+        let rope_emb = MiniCPMLongRoPE::new(&cfg, vb.device())?;
+        let lm_head = Linear::new(embed_tokens.embeddings().clone(), None);
+        Ok(Self {
+            cfg,
+            embed_tokens,
+            layers,
+            norm,
+            rope_emb,
+            lm_head
+        })
+    }
+
+    pub fn forward(&self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
+        let (bs, seq_len) = input_ids.dims2()?;
+        let input_embeds = self.embed_tokens.forward(&input_ids)?;
+        let attention_mask: Option<&Tensor> = {
+            if seq_len <= 1 {
+                None
+            } else {
+                Some(&prepare_causal_attention_mask(
+                    bs,
+                    seq_len,
+                    position_id,
+                    input_ids.device(),
+                )?)
+            }
+        };
+
+        let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
+        let mut hidden_states = input_embeds;
+        for decode_layer in &self.layers {
+            hidden_states = decode_layer.forward(&hidden_states, &cos, &sin, attention_mask)?;
+        }
+        hidden_states = self.norm.forward(&hidden_states)?;
+        let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?;
+        let logits = self.lm_head.forward(&hidden_state)?;
+        Ok(logits)
+    }
+
+    pub fn forward_step(&mut self, input_ids: &Tensor, position_id: usize) -> Result<Tensor> {
+        let (bs, seq_len) = input_ids.dims2()?;
+        let input_embeds = self.embed_tokens.forward(&input_ids)?;
+        let attention_mask: Option<&Tensor> = {
+            if seq_len <= 1 {
+                None
+            } else {
+                Some(&prepare_causal_attention_mask(
+                    bs,
+                    seq_len,
+                    position_id,
+                    input_ids.device(),
+                )?)
+            }
+        };
+
+        let (cos, sin) = self.rope_emb.forward(position_id, seq_len)?;
+        let mut hidden_states = input_embeds;
+        for decode_layer in &mut self.layers {
+            hidden_states = decode_layer.forward_step(&hidden_states, &cos, &sin, attention_mask)?;
+        }
+        hidden_states = self.norm.forward(&hidden_states)?;
+        let hidden_state = hidden_states.narrow(1, seq_len - 1, 1)?;
+        let logits = self.lm_head.forward(&hidden_state)?;
+        Ok(logits)
+    }
+
+    pub fn clear_kv_cache(&mut self) {
+        for layer in self.layers.iter_mut() {
+            layer.clear_kv_cache()
+        }
+    }
+}

src/models/mod.rs

@@ -1,4 +1,7 @@
 pub mod qwen2_5vl;
+pub mod minicpm4;
+pub mod base_modules;
+
 use anyhow::Result;
 use candle_core::{DType, Device};
 use openai_dive::v1::resources::chat::{

src/models/qwen2_5vl/config.rs

@@ -24,10 +24,10 @@ pub struct RopeScaling {
 }
 
 #[derive(Debug, Clone, PartialEq, serde::Deserialize)]
-pub struct Config {
+pub struct Qwen2_5VLConfig {
     pub attention_dropout: f32,
-    pub bos_token_id: usize,
-    pub eos_token_id: usize,
+    pub bos_token_id: u32,
+    pub eos_token_id: u32,
     pub vision_start_token_id: usize,
     pub vision_end_token_id: usize,
     pub vision_token_id: usize,

src/models/qwen2_5vl/generate.rs

@@ -1,5 +1,5 @@
 // use crate::models::GenerateStream;
-use crate::models::qwen2_5vl::config::Config;
+use crate::models::qwen2_5vl::config::Qwen2_5VLConfig;
 use crate::utils::utils::{
     build_completion_chunk_response, build_completion_response, find_safetensors_files, get_device,
     get_dtype, get_logit_processor,
@@ -15,7 +15,6 @@ use crate::{
 use anyhow::{Result, anyhow};
 use candle_core::{D, DType, Device, IndexOp, Tensor};
 use candle_nn::VarBuilder;
-use candle_transformers::generation::LogitsProcessor;
 use openai_dive::v1::resources::chat::{
     ChatCompletionChunkResponse, ChatCompletionParameters, ChatCompletionResponse,
 };
@@ -37,13 +36,13 @@ impl<'a> GenerateModel for Qwen2_5VLGenerateModel<'a> {
         let chat_template = ChatTemplate::init(path)?;
         let tokenizer = TokenizerModel::init(path)?;
         let config_path = path.to_string() + "/config.json";
-        let cfg: Config = serde_json::from_slice(&std::fs::read(config_path)?)?;
+        let cfg: Qwen2_5VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
         let device = &get_device(device);
         let cfg_dtype = cfg.torch_dtype.as_str();
         let dtype = get_dtype(dtype, cfg_dtype);
         let pre_processor = Qwen2_5VLProcessor::new(device, dtype)?;
-        let endoftext_id = cfg.bos_token_id as u32;
-        let im_end_id = cfg.eos_token_id as u32;
+        let endoftext_id = cfg.bos_token_id;
+        let im_end_id = cfg.eos_token_id;
         let model_list = find_safetensors_files(&path)?;
         let vb = unsafe { VarBuilder::from_mmaped_safetensors(&model_list, dtype, device)? };
         let qwen2_5_vl = Qwen2_5VLModel::new(cfg, vb)?;

src/models/qwen2_5vl/model.rs

@@ -1,12 +1,10 @@
 use crate::{
-    models::qwen2_5vl::config::{Config, RopeScaling},
+    models::qwen2_5vl::config::{Qwen2_5VLConfig, RopeScaling},
     position_embed::rope::{
-        Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding, apply_rotary_pos_emb,
-        apply_rotary_pos_emb_vision,
+        apply_rotary_pos_emb, apply_rotary_pos_emb_vision, Qwen2_5VLTextRotaryEmbedding, Qwen2_5VisionRotaryEmbedding
     },
     utils::tensor_utils::{
-        get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index,
-        safe_arg_sort_last_dim, zero_index,
+        get_equal_mask, get_vision_next_indices, masked_scatter_dim0, nonzero_index, repeat_kv, safe_arg_sort_last_dim, zero_index
     },
 };
 use anyhow::{Result, anyhow};
@@ -20,7 +18,7 @@ pub struct Qwen2_5VisionPatchEmbed {
 }
 
 impl Qwen2_5VisionPatchEmbed {
-    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let patch_size = cfg.vision_config.patch_size;
         let temporal_patch_size = cfg.vision_config.temporal_patch_size;
         let in_channels = cfg.vision_config.in_chans;
@@ -60,7 +58,7 @@ pub struct Qwen2_5VLPatchMerger {
 }
 
 impl Qwen2_5VLPatchMerger {
-    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let hidden_size =
             cfg.vision_config.hidden_size * (cfg.vision_config.spatial_merge_size.pow(2));
         let ln_q = rms_norm(
@@ -99,7 +97,7 @@ struct Qwen2_5VLVisionMLP {
 }
 
 impl Qwen2_5VLVisionMLP {
-    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let hidden_sz = cfg.vision_config.hidden_size;
         let intermediate_sz = cfg.vision_config.intermediate_size;
         let gate_proj = linear(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?;
@@ -131,7 +129,7 @@ struct Qwen2_5VLVisionAttention {
 }
 
 impl Qwen2_5VLVisionAttention {
-    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let hidden_size = cfg.vision_config.hidden_size;
         let num_heads = cfg.vision_config.num_heads;
         let head_dim = hidden_size / num_heads;
@@ -200,7 +198,7 @@ struct Qwen2_5VLVisionBlock {
 }
 
 impl Qwen2_5VLVisionBlock {
-    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let attn = Qwen2_5VLVisionAttention::new(cfg, vb.pp("attn"))?;
         let mlp = Qwen2_5VLVisionMLP::new(cfg, vb.pp("mlp"))?;
         let norm1 = rms_norm(
@@ -254,7 +252,7 @@ pub struct Qwen2_5VLVisionModel {
 }
 
 impl Qwen2_5VLVisionModel {
-    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let spatial_merge_size = cfg.vision_config.spatial_merge_size;
         let patch_size = cfg.vision_config.patch_size;
         let fullatt_block_indexes = cfg.vision_config.fullatt_block_indexes.clone();
@@ -539,23 +537,6 @@ impl Qwen2_5VLVisionModel {
     }
 }
 
-pub fn repeat_kv(xs: Tensor, n_rep: usize) -> Result<Tensor> {
-    if n_rep == 1 {
-        Ok(xs)
-    } else {
-        let (b_sz, n_kv_head, seq_len, head_dim) = xs.dims4()?;
-        // Using cat is faster than a broadcast as it avoids going through a potentially
-        // strided copy.
-        // https://github.com/huggingface/candle/pull/2043
-        let kv = Tensor::cat(&vec![&xs; n_rep], 2)?.reshape((
-            b_sz,
-            n_kv_head * n_rep,
-            seq_len,
-            head_dim,
-        ))?;
-        Ok(kv)
-    }
-}
 
 #[derive(Debug, Clone)]
 struct Qwen2_5VLTextMLP {
@@ -566,7 +547,7 @@ struct Qwen2_5VLTextMLP {
 }
 
 impl Qwen2_5VLTextMLP {
-    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let hidden_sz = cfg.hidden_size;
         let intermediate_sz = cfg.intermediate_size;
         let gate_proj = linear_no_bias(hidden_sz, intermediate_sz, vb.pp("gate_proj"))?;
@@ -605,7 +586,7 @@ struct Qwen2_5VLTextAttention {
 }
 
 impl Qwen2_5VLTextAttention {
-    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let hidden_size = cfg.hidden_size;
         let num_heads = cfg.num_attention_heads;
         let num_kv_heads = cfg.num_key_value_heads;
@@ -721,7 +702,7 @@ struct Qwen2_5VLTextDecoderLayer {
 }
 
 impl Qwen2_5VLTextDecoderLayer {
-    fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let self_attn = Qwen2_5VLTextAttention::new(cfg, vb.pp("self_attn"))?;
         let mlp = Qwen2_5VLTextMLP::new(cfg, vb.pp("mlp"))?;
         let input_layernorm =
@@ -774,7 +755,7 @@ pub struct Qwen2_5VLTextModel {
 }
 
 impl Qwen2_5VLTextModel {
-    pub fn new(cfg: &Config, vb: VarBuilder) -> Result<Self> {
+    pub fn new(cfg: &Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let embed_tokens =
             candle_nn::embedding(cfg.vocab_size, cfg.hidden_size, vb.pp("embed_tokens"))?;
         let head_dim = cfg.hidden_size / cfg.num_attention_heads;
@@ -879,13 +860,13 @@ impl Qwen2_5VLTextModel {
 pub struct Qwen2_5VLModel {
     visual: Qwen2_5VLVisionModel,
     model: Qwen2_5VLTextModel,
-    pub cfg: Config,
+    pub cfg: Qwen2_5VLConfig,
     lm_head: Linear,
     rope_deltas: Option<Tensor>,
 }
 
 impl Qwen2_5VLModel {
-    pub fn new(cfg: Config, vb: VarBuilder) -> Result<Self> {
+    pub fn new(cfg: Qwen2_5VLConfig, vb: VarBuilder) -> Result<Self> {
         let visual = Qwen2_5VLVisionModel::new(&cfg, vb.pp("visual"))?;
         let model = Qwen2_5VLTextModel::new(&cfg, vb.pp("model"))?;
         let vocab_size = cfg.vocab_size;

src/position_embed/rope.rs

@@ -2,8 +2,6 @@ use anyhow::Result;
 use candle_core::{D, DType, Device, IndexOp, Tensor};
 use candle_transformers::models::deepseek2::SplitOp;
 
-use crate::models::qwen2_5vl::config::RopeScaling;
-
 pub fn compute_default_rope_parameters(dim: usize, base: f32) -> Vec<f32> {
     let inv_freq: Vec<f32> = (0..dim)
         .step_by(2)

src/utils/tensor_utils.rs

@@ -1,6 +1,55 @@
 use anyhow::{Result, anyhow};
 use candle_core::{D, DType, Device, IndexOp, Tensor, shape::Dim};
 
+pub fn prepare_causal_attention_mask(
+    b_size: usize,
+    tgt_len: usize,
+    seqlen_offset: usize,
+    device: &Device
+) -> Result<Tensor> {
+    // Sliding window mask?
+    let mask: Vec<_> = (0..tgt_len)
+        .flat_map(|i| {
+            (0..tgt_len).map(move |j| {
+                if i < j {
+                    f32::NEG_INFINITY
+                } else {
+                    0.
+                }
+            })
+        })
+        .collect();
+    let mask = Tensor::from_slice(&mask, (tgt_len, tgt_len), device)?;
+    let mask = if seqlen_offset > 0 {
+        let mask0 = Tensor::zeros((tgt_len, seqlen_offset), DType::U32, device)?;
+        Tensor::cat(&[&mask0, &mask], D::Minus1)?
+    } else {
+        mask
+    };
+    let mask = mask
+        .expand((b_size, 1, tgt_len, tgt_len + seqlen_offset))?
+        .to_dtype(DType::U32)?;
+    Ok(mask)
+}
+
+pub fn repeat_kv(xs: Tensor, n_rep: usize) -> Result<Tensor> {
+    if n_rep == 1 {
+        Ok(xs)
+    } else {
+        let (b_sz, n_kv_head, seq_len, head_dim) = xs.dims4()?;
+        // Using cat is faster than a broadcast as it avoids going through a potentially
+        // strided copy.
+        // https://github.com/huggingface/candle/pull/2043
+        let kv = Tensor::cat(&vec![&xs; n_rep], 2)?.reshape((
+            b_sz,
+            n_kv_head * n_rep,
+            seq_len,
+            head_dim,
+        ))?;
+        Ok(kv)
+    }
+}
+
 pub fn split(t: &Tensor, splits: &[usize], dim: D) -> Result<Vec<Tensor>> {
     let dim = dim.to_index(t.shape(), "split")?;
     let mut split_res = Vec::new();

tests/config_tests.rs

@@ -1,12 +1,23 @@
-use aha::models::qwen2_5vl::config::Config;
+use aha::models::{minicpm4::config::MiniCPM4Config, qwen2_5vl::config::Qwen2_5VLConfig};
 use anyhow::Result;
 
 #[test]
-fn qwen2_5vl_config() -> Result<()> {
-    // cargo test qwen2_5vl_config -- --nocapture
+fn qwen2_5_vl_config() -> Result<()> {
+    // cargo test -F cuda,flash-attn qwen2_5vl_config -- --nocapture
     let model_path = "/home/jhq/huggingface_model/Qwen/Qwen2.5-VL-3B-Instruct/";
     let config_path = model_path.to_string() + "/config.json";
-    let config: Config = serde_json::from_slice(&std::fs::read(config_path)?)?;
+    let config: Qwen2_5VLConfig = serde_json::from_slice(&std::fs::read(config_path)?)?;
     println!("{:?}", config);
     Ok(())
 }
+
+
+#[test]
+fn minicpm4_config() -> Result<()> {
+    // cargo test -F cuda,flash-attn minicpm4_config -- --nocapture
+    let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/";
+    let config_path = model_path.to_string() + "/config.json";
+    let config: MiniCPM4Config = serde_json::from_slice(&std::fs::read(config_path)?)?;
+    println!("{:?}", config);
+    Ok(())
+}

tests/test_minicpm4.rs

@@ -0,0 +1,97 @@
+use std::time::Instant;
+
+use anyhow::Result;
+use candle_core::{DType, Device};
+use openai_dive::v1::resources::chat::ChatCompletionParameters;
+
+#[test]
+fn qwen2_5vl_generate() -> Result<()> {
+    // test with cpu :(太慢了, : RUST_BACKTRACE=1 cargo test qwen2_5vl_generate -- --nocapture
+    // test with cuda: RUST_BACKTRACE=1 cargo test -F cuda qwen2_5vl_generate -- --nocapture
+    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -- --nocapture
+    let device = Device::cuda_if_available(0)?;
+    let dtype = DType::BF16;
+
+    let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/";
+
+    let message = r#"
+    {
+        "model": "minicpm4",
+        "messages": [
+            {
+                "role": "user",
+                "content": [          
+                    {
+                        "type": "text", 
+                        "text": "你是谁"
+                    }
+                ]
+            }
+        ]
+    }
+    "#;
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    // let mut model = Qwen2_5VLGenerateModel::init(model_path, &device, dtype)?;
+    let mut model = ModelType::init(ModelType::Qwen2_5VL, model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+
+    let i_start = Instant::now();
+    let result = model.generate(mes)?;
+    println!("generate: \n {:?}", result);
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
+
+    Ok(())
+}
+
+#[tokio::test]
+async fn qwen2_5vl_stream() -> Result<()> {
+    // test with cuda+flash-attn: RUST_BACKTRACE=1 cargo test -F cuda,flash-attn qwen2_5vl_generate -- --nocapture
+    let device = Device::cuda_if_available(0)?;
+    let dtype = DType::BF16;
+
+    let model_path = "/home/jhq/huggingface_model/Qwen/Qwen2.5-VL-3B-Instruct/";
+
+    let message = r#"
+    {
+        "model": "qwen2.5vl",
+        "messages": [
+            {
+                "role": "user",
+                "content": [ 
+                    {
+                        "type": "image",
+                        "image_url": 
+                        {
+                            "url": "file://./assets/img/ocr_test.png"
+                        }
+                    },               
+                    {
+                        "type": "text", 
+                        "text": "请分析图片并提取所有可见文本内容，按从左到右、从上到下的布局，返回纯文本"
+                    }
+                ]
+            }
+        ]
+    }
+    "#;
+    let mes: ChatCompletionParameters = serde_json::from_str(message)?;
+    let i_start = Instant::now();
+    // let mut model = Qwen2_5VLGenerateModel::init(model_path, &device, dtype)?;
+    let mut model = ModelType::init(ModelType::Qwen2_5VL, model_path, None, None)?;
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in load model is: {:?}", i_duration);
+
+    let i_start = Instant::now();
+    let mut stream = pin!(model.generate_stream(mes)?);
+    while let Some(item) = stream.next().await {
+        println!("generate: \n {:?}", item);
+    }
+
+    let i_duration = i_start.elapsed();
+    println!("Time elapsed in generate is: {:?}", i_duration);
+
+    Ok(())
+}

tests/weight_test.rs

@@ -0,0 +1,18 @@
+use aha::utils::utils::find_safetensors_files;
+use anyhow::Result;
+use candle_core::{safetensors, Device};
+#[test]
+fn minicpm4_weight() -> Result<()> {
+    let model_path = "/home/jhq/huggingface_model/OpenBMB/MiniCPM4-0.5B/";
+    let model_list = find_safetensors_files(&model_path)?;
+    let device = Device::Cpu;
+    for m in model_list {
+        let weights = safetensors::load(m, &device)?;
+        for (key, tensor) in weights.iter() {
+            println!("=== {} ===", key);
+            println!("Shape: {:?}", tensor.shape());
+            println!("DType: {:?}", tensor.dtype());
+        }
+    }
+    Ok(())
+}