refactor(项目结构) 重构为 workspace 多成员项目

- 将项目拆分为三个 crate：minicpm-core（核心模型）、minicpm-convert（转换功能）、minicpm-inference（推理功能）
- 添加两个示例：minimal-inference（最小推理）和 convert（模型转换）
- 转换后自动拷贝 config.json 和 tokenizer.json 到 model 目录
- 更新 README 说明 workspace 结构和使用方式

crates/minicpm-convert/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "minicpm-convert"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+minicpm-core = { path = "../minicpm-core" }
+burn = { version = "0.21", features = ["std", "wgpu"] }
+memmap2 = "0.9"
+anyhow = "1.0"
+serde_json = "1.0"

crates/minicpm-convert/src/lib.rs

@@ -0,0 +1,258 @@
+use minicpm_core::config::LlamaConfig;
+use minicpm_core::model::LlamaForCausalLM;
+use burn::module::{Module, Param};
+use burn::nn::{EmbeddingRecord, LinearRecord};
+use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Shape, Tensor, TensorData};
+use memmap2::Mmap;
+use std::collections::HashMap;
+use std::path::Path;
+
+pub fn export_model<B: Backend>(
+    safetensors_path: &Path,
+    config_path: &Path,
+    tokenizer_path: &Path,
+    output_dir: &Path,
+    device: &B::Device,
+) -> anyhow::Result<()> {
+    println!("开始转换 MiniCPM 模型为 Burn 格式...");
+
+    println!("加载配置文件: {:?}", config_path);
+    let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;
+
+    println!("创建模型结构...");
+    let model = LlamaForCausalLM::<B>::new(config, device);
+
+    println!("加载 safetensors 权重...");
+    let model = load_safetensors(model, safetensors_path, device)?;
+
+    println!("保存为 MPK 格式...");
+    std::fs::create_dir_all(output_dir)?;
+    let output_path = output_dir.join("model");
+    let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
+    model.save_file(&output_path, &recorder)?;
+
+    println!("拷贝配置文件和 tokenizer...");
+    std::fs::copy(config_path, output_dir.join("config.json"))?;
+    std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
+
+    println!("模型已成功导出到: {:?}", output_dir);
+    Ok(())
+}
+pub fn load_safetensors<B: Backend>(
+    model: LlamaForCausalLM<B>,
+    path: &Path,
+    device: &B::Device,
+) -> anyhow::Result<LlamaForCausalLM<B>> {
+    let file = std::fs::File::open(path)?;
+    let mmap = unsafe { Mmap::map(&file)? };
+
+    let header_len = u64::from_le_bytes(mmap[..8].try_into().unwrap()) as usize;
+    let header_bytes = &mmap[8..8 + header_len];
+    let header: serde_json::Value = serde_json::from_slice(header_bytes)?;
+
+    let mut tensors = HashMap::new();
+    let mut shapes = HashMap::new();
+    let mut dtypes = HashMap::new();
+
+    let data_offset = 8 + header_len;
+
+    if let Some(obj) = header.as_object() {
+        for (name, info) in obj {
+            if name == "__metadata__" {
+                continue;
+            }
+            if let Some(info_obj) = info.as_object() {
+                if let Some(offsets) = info_obj.get("data_offsets") {
+                    if let Some(arr) = offsets.as_array() {
+                        let start = arr[0].as_u64().unwrap() as usize;
+                        let end = arr[1].as_u64().unwrap() as usize;
+                        let data = &mmap[data_offset + start..data_offset + end];
+                        tensors.insert(name.clone(), data.to_vec());
+                    }
+                }
+                if let Some(shape) = info_obj.get("shape") {
+                    if let Some(arr) = shape.as_array() {
+                        let shape_vec: Vec<usize> = arr
+                            .iter()
+                            .map(|v| v.as_u64().unwrap() as usize)
+                            .collect();
+                        shapes.insert(name.clone(), shape_vec);
+                    }
+                }
+                if let Some(dtype) = info_obj.get("dtype") {
+                    if let Some(s) = dtype.as_str() {
+                        dtypes.insert(name.clone(), s.to_string());
+                    }
+                }
+            }
+        }
+    }
+
+    let mut model = model;
+
+    if let Some(weight) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, "model.embed_tokens.weight", device) {
+        let record = EmbeddingRecord {
+            weight: Param::from_tensor(weight),
+        };
+        model.model.embed_tokens = model.model.embed_tokens.clone().load_record(record);
+    }
+
+    for i in 0..model.config.num_hidden_layers {
+        let prefix = format!("model.layers.{i}");
+
+        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.q_proj.weight"), device) {
+            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+            model.model.layers[i].self_attn.q_proj = model.model.layers[i].self_attn.q_proj.clone().load_record(record);
+        }
+        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.k_proj.weight"), device) {
+            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+            model.model.layers[i].self_attn.k_proj = model.model.layers[i].self_attn.k_proj.clone().load_record(record);
+        }
+        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.v_proj.weight"), device) {
+            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+            model.model.layers[i].self_attn.v_proj = model.model.layers[i].self_attn.v_proj.clone().load_record(record);
+        }
+        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.self_attn.o_proj.weight"), device) {
+            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+            model.model.layers[i].self_attn.o_proj = model.model.layers[i].self_attn.o_proj.clone().load_record(record);
+        }
+
+        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.gate_proj.weight"), device) {
+            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+            model.model.layers[i].mlp.gate_proj = model.model.layers[i].mlp.gate_proj.clone().load_record(record);
+        }
+        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.up_proj.weight"), device) {
+            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+            model.model.layers[i].mlp.up_proj = model.model.layers[i].mlp.up_proj.clone().load_record(record);
+        }
+        if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.mlp.down_proj.weight"), device) {
+            let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+            model.model.layers[i].mlp.down_proj = model.model.layers[i].mlp.down_proj.clone().load_record(record);
+        }
+
+        if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.input_layernorm.weight"), device) {
+            model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
+        }
+        if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, &format!("{prefix}.post_attention_layernorm.weight"), device) {
+            model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
+        }
+    }
+
+    if let Some(w) = load_tensor_1d::<B>(&tensors, &shapes, &dtypes, "model.norm.weight", device) {
+        model.model.norm.weight = Param::from_tensor(w);
+    }
+
+    if let Some(w) = load_tensor_2d::<B>(&tensors, &shapes, &dtypes, "lm_head.weight", device) {
+        let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+        model.lm_head = model.lm_head.clone().load_record(record);
+    } else if model.config.tie_word_embeddings {
+        let embed_weight = model.model.embed_tokens.weight.val();
+        let record = LinearRecord { weight: Param::from_tensor(embed_weight.transpose()), bias: None };
+        model.lm_head = model.lm_head.clone().load_record(record);
+    }
+
+    Ok(model)
+}
+
+fn load_tensor_2d<B: Backend>(
+    tensors: &HashMap<String, Vec<u8>>,
+    shapes: &HashMap<String, Vec<usize>>,
+    dtypes: &HashMap<String, String>,
+    name: &str,
+    device: &B::Device,
+) -> Option<Tensor<B, 2, Float>> {
+    let data = tensors.get(name)?;
+    let shape = shapes.get(name)?;
+    let dtype = dtypes.get(name)?;
+
+    assert_eq!(shape.len(), 2, "Expected 2D tensor for {name}");
+
+    let f32_data = convert_to_f32(data, dtype);
+    let shape_arr: [usize; 2] = [shape[0], shape[1]];
+    let tensor_data = TensorData::new(f32_data, Shape::new(shape_arr));
+    Some(Tensor::from_data(tensor_data, device))
+}
+
+fn load_tensor_1d<B: Backend>(
+    tensors: &HashMap<String, Vec<u8>>,
+    shapes: &HashMap<String, Vec<usize>>,
+    dtypes: &HashMap<String, String>,
+    name: &str,
+    device: &B::Device,
+) -> Option<Tensor<B, 1, Float>> {
+    let data = tensors.get(name)?;
+    let shape = shapes.get(name)?;
+    let dtype = dtypes.get(name)?;
+
+    assert_eq!(shape.len(), 1, "Expected 1D tensor for {name}");
+
+    let f32_data = convert_to_f32(data, dtype);
+    let shape_arr: [usize; 1] = [shape[0]];
+    let tensor_data = TensorData::new(f32_data, Shape::new(shape_arr));
+    Some(Tensor::from_data(tensor_data, device))
+}
+
+fn convert_to_f32(data: &[u8], dtype: &str) -> Vec<f32> {
+    match dtype {
+        "BF16" => data
+            .chunks_exact(2)
+            .map(|chunk| {
+                let bytes: [u8; 2] = chunk.try_into().unwrap();
+                bf16_to_f32(u16::from_le_bytes(bytes))
+            })
+            .collect(),
+        "F16" => data
+            .chunks_exact(2)
+            .map(|chunk| {
+                let bytes: [u8; 2] = chunk.try_into().unwrap();
+                f16_to_f32(u16::from_le_bytes(bytes))
+            })
+            .collect(),
+        "F32" => data
+            .chunks_exact(4)
+            .map(|chunk| {
+                let bytes: [u8; 4] = chunk.try_into().unwrap();
+                f32::from_le_bytes(bytes)
+            })
+            .collect(),
+        _ => panic!("Unsupported dtype: {dtype}"),
+    }
+}
+
+fn bf16_to_f32(bf16: u16) -> f32 {
+    let f32_bits = (bf16 as u32) << 16;
+    f32::from_bits(f32_bits)
+}
+
+fn f16_to_f32(f16: u16) -> f32 {
+    let sign = (f16 >> 15) & 1;
+    let exp = (f16 >> 10) & 0x1F;
+    let mant = f16 & 0x3FF;
+
+    if exp == 0 {
+        if mant == 0 {
+            return if sign == 0 { 0.0 } else { -0.0 };
+        }
+        let mut m = mant as u32;
+        let mut e = 0u32;
+        while (m & 0x400) == 0 {
+            m <<= 1;
+            e += 1;
+        }
+        m &= 0x3FF;
+        let f32_exp = 127 - 14 - e;
+        let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((m as u32) << 13);
+        return f32::from_bits(f32_bits);
+    }
+
+    if exp == 0x1F {
+        let f32_bits = ((sign as u32) << 31) | (0xFFu32 << 23) | ((mant as u32) << 13);
+        return f32::from_bits(f32_bits);
+    }
+
+    let f32_exp = (exp as u32) + 127 - 15;
+    let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((mant as u32) << 13);
+    f32::from_bits(f32_bits)
+}

crates/minicpm-core/Cargo.toml

@@ -0,0 +1,10 @@
+[package]
+name = "minicpm-core"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+burn = { version = "0.21", features = ["std", "wgpu"] }
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"
+anyhow = "1.0"

crates/minicpm-core/src/config.rs

@@ -0,0 +1,62 @@
+use serde::Deserialize;
+
+#[derive(Debug, Clone, Deserialize)]
+pub struct LlamaConfig {
+    pub hidden_size: usize,
+    pub intermediate_size: usize,
+    pub num_hidden_layers: usize,
+    pub num_attention_heads: usize,
+    pub num_key_value_heads: usize,
+    pub vocab_size: usize,
+    pub hidden_act: String,
+    pub rms_norm_eps: f64,
+    pub rope_theta: f64,
+    pub max_position_embeddings: usize,
+    pub bos_token_id: u32,
+    pub eos_token_id: EosTokenId,
+    pub pad_token_id: u32,
+    pub tie_word_embeddings: bool,
+    pub head_dim: Option<usize>,
+}
+
+#[derive(Debug, Clone, Deserialize)]
+#[serde(untagged)]
+pub enum EosTokenId {
+    Single(u32),
+    Multiple(Vec<u32>),
+}
+
+impl EosTokenId {
+    pub fn first(&self) -> u32 {
+        match self {
+            EosTokenId::Single(id) => *id,
+            EosTokenId::Multiple(ids) => ids[0],
+        }
+    }
+
+    pub fn contains(&self, token: u32) -> bool {
+        match self {
+            EosTokenId::Single(id) => *id == token,
+            EosTokenId::Multiple(ids) => ids.contains(&token),
+        }
+    }
+
+    pub fn as_slice(&self) -> &[u32] {
+        match self {
+            EosTokenId::Single(id) => std::slice::from_ref(id),
+            EosTokenId::Multiple(ids) => ids.as_slice(),
+        }
+    }
+}
+
+impl LlamaConfig {
+    pub fn head_dim(&self) -> usize {
+        self.head_dim.unwrap_or(self.hidden_size / self.num_attention_heads)
+    }
+
+    pub fn from_json(path: &str) -> anyhow::Result<Self> {
+        let content = std::fs::read_to_string(path)?;
+        let config: LlamaConfig = serde_json::from_str(&content)?;
+        Ok(config)
+    }
+}

crates/minicpm-core/src/lib.rs

@@ -0,0 +1,5 @@
+pub mod config;
+pub mod model;
+
+pub use config::{EosTokenId, LlamaConfig};
+pub use model::{LlamaForCausalLM, LlamaKVCache};

crates/minicpm-core/src/model/attention.rs

@@ -0,0 +1,186 @@
+use burn::module::Module;
+use burn::nn::{Linear, LinearConfig};
+use burn::tensor::activation::softmax;
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Tensor};
+
+use super::rope::RoPE;
+
+#[derive(Debug, Clone)]
+pub struct KVCache<B: Backend> {
+    pub k: Tensor<B, 4, Float>,
+    pub v: Tensor<B, 4, Float>,
+}
+
+#[derive(Module, Debug)]
+pub struct Attention<B: Backend> {
+    pub q_proj: Linear<B>,
+    pub k_proj: Linear<B>,
+    pub v_proj: Linear<B>,
+    pub o_proj: Linear<B>,
+    pub n_heads: usize,
+    pub n_kv_heads: usize,
+    pub head_dim: usize,
+}
+
+impl<B: Backend> Attention<B> {
+    pub fn new(
+        hidden_size: usize,
+        n_heads: usize,
+        n_kv_heads: usize,
+        head_dim: usize,
+        device: &B::Device,
+    ) -> Self {
+        let q_proj = LinearConfig::new(hidden_size, n_heads * head_dim).init(device);
+        let k_proj = LinearConfig::new(hidden_size, n_kv_heads * head_dim).init(device);
+        let v_proj = LinearConfig::new(hidden_size, n_kv_heads * head_dim).init(device);
+        let o_proj = LinearConfig::new(n_heads * head_dim, hidden_size).init(device);
+
+        Self {
+            q_proj,
+            k_proj,
+            v_proj,
+            o_proj,
+            n_heads,
+            n_kv_heads,
+            head_dim,
+        }
+    }
+
+    pub fn forward(
+        &self,
+        x: Tensor<B, 3, Float>,
+        rope: &RoPE,
+        mask: Option<Tensor<B, 4, Float>>,
+        offset: usize,
+    ) -> Tensor<B, 3, Float> {
+        let shape = x.shape();
+        let dims: [usize; 3] = shape.dims();
+        let batch = dims[0];
+        let seq_len = dims[1];
+
+        let q = self.q_proj.forward(x.clone());
+        let k = self.k_proj.forward(x.clone());
+        let v = self.v_proj.forward(x);
+
+        let q = q
+            .reshape([batch, seq_len, self.n_heads, self.head_dim])
+            .swap_dims(1, 2);
+        let k = k
+            .reshape([batch, seq_len, self.n_kv_heads, self.head_dim])
+            .swap_dims(1, 2);
+        let v = v
+            .reshape([batch, seq_len, self.n_kv_heads, self.head_dim])
+            .swap_dims(1, 2);
+
+        let q = rope.apply(q, offset);
+        let k = rope.apply(k, offset);
+
+        let k = self.repeat_kv(k);
+        let v = self.repeat_kv(v);
+
+        let scale = (self.head_dim as f64).sqrt() as f32;
+        let scores = q
+            .matmul(k.swap_dims(2, 3))
+            .div_scalar(scale);
+
+        let scores = match mask {
+            Some(mask) => scores + mask,
+            None => scores,
+        };
+
+        let attn = softmax(scores, 3);
+        let output = attn.matmul(v);
+
+        let output = output
+            .swap_dims(1, 2)
+            .reshape([batch, seq_len, self.n_heads * self.head_dim]);
+
+        self.o_proj.forward(output)
+    }
+
+    pub fn forward_with_cache(
+        &self,
+        x: Tensor<B, 3, Float>,
+        rope: &RoPE,
+        mask: Option<Tensor<B, 4, Float>>,
+        offset: usize,
+        cache: Option<&KVCache<B>>,
+    ) -> (Tensor<B, 3, Float>, KVCache<B>) {
+        let shape = x.shape();
+        let dims: [usize; 3] = shape.dims();
+        let batch = dims[0];
+        let seq_len = dims[1];
+
+        let q = self.q_proj.forward(x.clone());
+        let k = self.k_proj.forward(x.clone());
+        let v = self.v_proj.forward(x);
+
+        let q = q
+            .reshape([batch, seq_len, self.n_heads, self.head_dim])
+            .swap_dims(1, 2);
+        let k = k
+            .reshape([batch, seq_len, self.n_kv_heads, self.head_dim])
+            .swap_dims(1, 2);
+        let v = v
+            .reshape([batch, seq_len, self.n_kv_heads, self.head_dim])
+            .swap_dims(1, 2);
+
+        let q = rope.apply(q, offset);
+        let k = rope.apply(k, offset);
+
+        // 合并 KV Cache
+        let (k_full, v_full) = match cache {
+            Some(cache) => {
+                let k_full = Tensor::cat(vec![cache.k.clone(), k], 2);
+                let v_full = Tensor::cat(vec![cache.v.clone(), v], 2);
+                (k_full, v_full)
+            }
+            None => (k.clone(), v.clone()),
+        };
+
+        let new_cache = KVCache {
+            k: k_full.clone(),
+            v: v_full.clone(),
+        };
+
+        let k = self.repeat_kv(k_full);
+        let v = self.repeat_kv(v_full);
+
+        let scale = (self.head_dim as f64).sqrt() as f32;
+        let scores = q
+            .matmul(k.swap_dims(2, 3))
+            .div_scalar(scale);
+
+        let scores = match mask {
+            Some(mask) => scores + mask,
+            None => scores,
+        };
+
+        let attn = softmax(scores, 3);
+        let output = attn.matmul(v);
+
+        let output = output
+            .swap_dims(1, 2)
+            .reshape([batch, seq_len, self.n_heads * self.head_dim]);
+
+        (self.o_proj.forward(output), new_cache)
+    }
+
+    fn repeat_kv(&self, x: Tensor<B, 4, Float>) -> Tensor<B, 4, Float> {
+        let n_rep = self.n_heads / self.n_kv_heads;
+        if n_rep == 1 {
+            return x;
+        }
+
+        let shape = x.shape();
+        let dims: [usize; 4] = shape.dims();
+        let batch = dims[0];
+        let kv_heads = dims[1];
+        let seq_len = dims[2];
+        let head_dim = dims[3];
+
+        let x = x.reshape([batch, kv_heads, 1, seq_len, head_dim]).expand([batch, kv_heads, n_rep, seq_len, head_dim]);
+        x.reshape([batch, kv_heads * n_rep, seq_len, head_dim])
+    }
+}

crates/minicpm-core/src/model/decoder.rs

@@ -0,0 +1,81 @@
+use burn::module::Module;
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Tensor};
+
+use super::attention::{Attention, KVCache};
+use super::ffn::FeedForward;
+use super::norm::RmsNorm;
+use super::rope::RoPE;
+
+#[derive(Module, Debug)]
+pub struct DecoderLayer<B: Backend> {
+    pub self_attn: Attention<B>,
+    pub mlp: FeedForward<B>,
+    pub input_layernorm: RmsNorm<B>,
+    pub post_attention_layernorm: RmsNorm<B>,
+}
+
+impl<B: Backend> DecoderLayer<B> {
+    pub fn new(
+        hidden_size: usize,
+        intermediate_size: usize,
+        n_heads: usize,
+        n_kv_heads: usize,
+        head_dim: usize,
+        rms_norm_eps: f64,
+        device: &B::Device,
+    ) -> Self {
+        let self_attn = Attention::new(hidden_size, n_heads, n_kv_heads, head_dim, device);
+        let mlp = FeedForward::new(hidden_size, intermediate_size, device);
+        let input_layernorm = RmsNorm::new(hidden_size, rms_norm_eps, device);
+        let post_attention_layernorm = RmsNorm::new(hidden_size, rms_norm_eps, device);
+
+        Self {
+            self_attn,
+            mlp,
+            input_layernorm,
+            post_attention_layernorm,
+        }
+    }
+
+    pub fn forward(
+        &self,
+        x: Tensor<B, 3, Float>,
+        rope: &RoPE,
+        mask: Option<Tensor<B, 4, Float>>,
+        offset: usize,
+    ) -> Tensor<B, 3, Float> {
+        let residual = x.clone();
+        let x = self.input_layernorm.forward(x);
+        let x = self.self_attn.forward(x, rope, mask, offset);
+        let x = residual + x;
+
+        let residual = x.clone();
+        let x = self.post_attention_layernorm.forward(x);
+        let x = self.mlp.forward(x);
+        let x = residual + x;
+
+        x
+    }
+
+    pub fn forward_with_cache(
+        &self,
+        x: Tensor<B, 3, Float>,
+        rope: &RoPE,
+        mask: Option<Tensor<B, 4, Float>>,
+        offset: usize,
+        cache: Option<&KVCache<B>>,
+    ) -> (Tensor<B, 3, Float>, KVCache<B>) {
+        let residual = x.clone();
+        let x = self.input_layernorm.forward(x);
+        let (x, kv_cache) = self.self_attn.forward_with_cache(x, rope, mask, offset, cache);
+        let x = residual + x;
+
+        let residual = x.clone();
+        let x = self.post_attention_layernorm.forward(x);
+        let x = self.mlp.forward(x);
+        let x = residual + x;
+
+        (x, kv_cache)
+    }
+}

crates/minicpm-core/src/model/ffn.rs

@@ -0,0 +1,33 @@
+use burn::module::Module;
+use burn::nn::{Linear, LinearConfig};
+use burn::tensor::backend::Backend;
+use burn::tensor::activation::silu;
+use burn::tensor::{Float, Tensor};
+
+#[derive(Module, Debug)]
+pub struct FeedForward<B: Backend> {
+    pub gate_proj: Linear<B>,
+    pub up_proj: Linear<B>,
+    pub down_proj: Linear<B>,
+}
+
+impl<B: Backend> FeedForward<B> {
+    pub fn new(hidden_size: usize, intermediate_size: usize, device: &B::Device) -> Self {
+        let gate_proj = LinearConfig::new(hidden_size, intermediate_size).init(device);
+        let up_proj = LinearConfig::new(hidden_size, intermediate_size).init(device);
+        let down_proj = LinearConfig::new(intermediate_size, hidden_size).init(device);
+
+        Self {
+            gate_proj,
+            up_proj,
+            down_proj,
+        }
+    }
+
+    pub fn forward(&self, x: Tensor<B, 3, Float>) -> Tensor<B, 3, Float> {
+        let gate = self.gate_proj.forward(x.clone());
+        let up = self.up_proj.forward(x);
+        let activated = silu(gate) * up;
+        self.down_proj.forward(activated)
+    }
+}

crates/minicpm-core/src/model/mod.rs

@@ -0,0 +1,8 @@
+pub mod attention;
+pub mod decoder;
+pub mod ffn;
+pub mod model;
+pub mod norm;
+pub mod rope;
+
+pub use model::{LlamaForCausalLM, LlamaKVCache};

crates/minicpm-core/src/model/model.rs

@@ -0,0 +1,164 @@
+use burn::module::Module;
+use burn::nn::{Embedding, EmbeddingConfig, Linear, LinearConfig};
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Int, Tensor};
+
+use super::attention::KVCache;
+use crate::config::LlamaConfig;
+use super::decoder::DecoderLayer;
+use super::norm::RmsNorm;
+use super::rope::RoPE;
+
+#[derive(Debug, Clone)]
+pub struct LlamaKVCache<B: Backend> {
+    pub layers: Vec<KVCache<B>>,
+}
+
+impl<B: Backend> LlamaKVCache<B> {
+    pub fn new(num_layers: usize) -> Self {
+        Self {
+            layers: Vec::with_capacity(num_layers),
+        }
+    }
+}
+
+#[derive(Module, Debug)]
+pub struct LlamaModel<B: Backend> {
+    pub embed_tokens: Embedding<B>,
+    pub layers: Vec<DecoderLayer<B>>,
+    pub norm: RmsNorm<B>,
+    pub rope: RoPE,
+}
+
+#[derive(Module, Debug)]
+pub struct LlamaForCausalLM<B: Backend> {
+    pub model: LlamaModel<B>,
+    pub lm_head: Linear<B>,
+    pub config: LlamaConfig,
+}
+
+impl<B: Backend> LlamaForCausalLM<B> {
+    pub fn new(config: LlamaConfig, device: &B::Device) -> Self {
+        let head_dim = config.head_dim();
+        let rope = RoPE::new(head_dim, config.rope_theta, config.max_position_embeddings);
+
+        let embed_tokens =
+            EmbeddingConfig::new(config.vocab_size, config.hidden_size).init(device);
+
+        let mut layers = Vec::with_capacity(config.num_hidden_layers);
+        for _ in 0..config.num_hidden_layers {
+            layers.push(DecoderLayer::new(
+                config.hidden_size,
+                config.intermediate_size,
+                config.num_attention_heads,
+                config.num_key_value_heads,
+                head_dim,
+                config.rms_norm_eps,
+                device,
+            ));
+        }
+
+        let norm = RmsNorm::new(config.hidden_size, config.rms_norm_eps, device);
+
+        let lm_head = LinearConfig::new(config.hidden_size, config.vocab_size).init(device);
+
+        let model = LlamaModel {
+            embed_tokens,
+            layers,
+            norm,
+            rope,
+        };
+
+        Self {
+            model,
+            lm_head,
+            config,
+        }
+    }
+
+    pub fn forward(&self, input_ids: Tensor<B, 2, Int>) -> Tensor<B, 3, Float> {
+        let shape = input_ids.shape();
+        let dims: [usize; 2] = shape.dims();
+        let seq_len = dims[1];
+
+        let mask = if seq_len > 1 {
+            Some(self.make_causal_mask(seq_len, &input_ids.device()))
+        } else {
+            None
+        };
+
+        let x = self.model.embed_tokens.forward(input_ids);
+
+        let mut x = x;
+        for layer in &self.model.layers {
+            x = layer.forward(x, &self.model.rope, mask.clone(), 0);
+        }
+
+        let x = self.model.norm.forward(x);
+
+        let logits = self.lm_head.forward(x);
+
+        logits
+    }
+
+    pub fn forward_with_cache(
+        &self,
+        input_ids: Tensor<B, 2, Int>,
+        cache: Option<&LlamaKVCache<B>>,
+    ) -> (Tensor<B, 3, Float>, LlamaKVCache<B>) {
+        let shape = input_ids.shape();
+        let dims: [usize; 2] = shape.dims();
+        let seq_len = dims[1];
+
+        let offset = match &cache {
+            Some(cache) => {
+                let k_dims: [usize; 4] = cache.layers[0].k.shape().dims();
+                k_dims[2]
+            }
+            None => 0,
+        };
+
+        let mask = if seq_len > 1 {
+            Some(self.make_causal_mask(seq_len, &input_ids.device()))
+        } else {
+            None
+        };
+
+        let x = self.model.embed_tokens.forward(input_ids);
+
+        let mut x = x;
+        let mut new_cache = LlamaKVCache::new(self.config.num_hidden_layers);
+
+        for (i, layer) in self.model.layers.iter().enumerate() {
+            let layer_cache = cache.as_ref().map(|c| &c.layers[i]);
+            let (output, kv_cache) = layer.forward_with_cache(
+                x,
+                &self.model.rope,
+                mask.clone(),
+                offset,
+                layer_cache,
+            );
+            x = output;
+            new_cache.layers.push(kv_cache);
+        }
+
+        let x = self.model.norm.forward(x);
+
+        let logits = self.lm_head.forward(x);
+
+        (logits, new_cache)
+    }
+
+    fn make_causal_mask(&self, seq_len: usize, device: &B::Device) -> Tensor<B, 4, Float> {
+        let mut mask = vec![0.0f32; seq_len * seq_len];
+        for i in 0..seq_len {
+            for j in 0..seq_len {
+                if j > i {
+                    mask[i * seq_len + j] = f32::NEG_INFINITY;
+                }
+            }
+        }
+        Tensor::<B, 1, Float>::from_floats(mask.as_slice(), device)
+            .reshape([1, 1, seq_len, seq_len])
+    }
+}

crates/minicpm-core/src/model/norm.rs

@@ -0,0 +1,29 @@
+use burn::module::{Module, Param};
+use burn::nn::Initializer;
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Tensor};
+
+#[derive(Module, Debug)]
+pub struct RmsNorm<B: Backend> {
+    pub weight: Param<Tensor<B, 1, Float>>,
+    pub eps: f64,
+}
+
+impl<B: Backend> RmsNorm<B> {
+    pub fn new(dim: usize, eps: f64, device: &B::Device) -> Self {
+        let weight = Initializer::Ones.init([dim], device);
+        Self {
+            weight,
+            eps,
+        }
+    }
+
+    pub fn forward<const D: usize>(&self, x: Tensor<B, D, Float>) -> Tensor<B, D, Float> {
+        let variance = x.clone().powf_scalar(2.0).mean_dim(D - 1);
+        let rms = (variance + self.eps).sqrt();
+        let x_norm = x / rms;
+
+        let weight = self.weight.val().clone().unsqueeze();
+        x_norm * weight
+    }
+}

crates/minicpm-core/src/model/rope.rs

@@ -0,0 +1,83 @@
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Tensor};
+
+#[derive(Debug, Clone)]
+pub struct RoPE {
+    pub dim: usize,
+    pub theta: f64,
+    pub max_seq_len: usize,
+    cos_cache: Vec<f32>,
+    sin_cache: Vec<f32>,
+    half_dim: usize,
+}
+
+impl RoPE {
+    pub fn new(dim: usize, theta: f64, max_seq_len: usize) -> Self {
+        let half_dim = dim / 2;
+
+        let mut inv_freqs = Vec::with_capacity(half_dim);
+        for i in 0..half_dim {
+            let freq = 1.0 / (theta.powf((2.0 * i as f64) / dim as f64));
+            inv_freqs.push(freq as f32);
+        }
+
+        let mut cos_cache = Vec::with_capacity(max_seq_len * half_dim);
+        let mut sin_cache = Vec::with_capacity(max_seq_len * half_dim);
+
+        for pos in 0..max_seq_len {
+            for &freq in &inv_freqs {
+                let angle = pos as f32 * freq;
+                cos_cache.push(angle.cos());
+                sin_cache.push(angle.sin());
+            }
+        }
+
+        Self {
+            dim,
+            theta,
+            max_seq_len,
+            cos_cache,
+            sin_cache,
+            half_dim,
+        }
+    }
+
+    pub fn get_cache(&self, offset: usize, len: usize) -> (Vec<f32>, Vec<f32>) {
+        let cos = self.cos_cache[offset * self.half_dim..(offset + len) * self.half_dim].to_vec();
+        let sin = self.sin_cache[offset * self.half_dim..(offset + len) * self.half_dim].to_vec();
+        (cos, sin)
+    }
+
+    pub fn apply<B: Backend>(
+        &self,
+        x: Tensor<B, 4, Float>,
+        offset: usize,
+    ) -> Tensor<B, 4, Float> {
+        let device = x.device();
+        let shape = x.shape();
+        let dims: [usize; 4] = shape.dims();
+        let batch = dims[0];
+        let heads = dims[1];
+        let seq_len = dims[2];
+        let dim = dims[3];
+        let half_dim = self.half_dim;
+
+        let cos_vals = &self.cos_cache[offset * half_dim..(offset + seq_len) * half_dim];
+        let sin_vals = &self.sin_cache[offset * half_dim..(offset + seq_len) * half_dim];
+
+        let cos = Tensor::<B, 1, Float>::from_floats(cos_vals, &device)
+            .reshape([1, 1, seq_len, half_dim])
+            .expand([batch, heads, seq_len, half_dim]);
+        let sin = Tensor::<B, 1, Float>::from_floats(sin_vals, &device)
+            .reshape([1, 1, seq_len, half_dim])
+            .expand([batch, heads, seq_len, half_dim]);
+
+        let x1 = x.clone().slice([0..batch, 0..heads, 0..seq_len, 0..half_dim]);
+        let x2 = x.slice([0..batch, 0..heads, 0..seq_len, half_dim..dim]);
+
+        let x1_rot = x1.clone() * cos.clone() - x2.clone() * sin.clone();
+        let x2_rot = x2 * cos + x1 * sin;
+
+        Tensor::cat(vec![x1_rot, x2_rot], 3)
+    }
+}

crates/minicpm-inference/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "minicpm-inference"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+minicpm-core = { path = "../minicpm-core" }
+burn = { version = "0.21", features = ["std", "wgpu"] }
+tokenizers = "0.20"
+rand = "0.8"
+anyhow = "1.0"

crates/minicpm-inference/src/lib.rs

@@ -0,0 +1,283 @@
+pub use minicpm_core::config::{EosTokenId, LlamaConfig};
+pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
+
+use burn::module::Module;
+use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Int, Tensor};
+use rand::Rng;
+use tokenizers::Tokenizer;
+
+
+pub struct TokenizerWrapper {
+    tokenizer: Tokenizer,
+}
+
+impl TokenizerWrapper {
+    pub fn from_file(path: &str) -> anyhow::Result<Self> {
+        let tokenizer = Tokenizer::from_file(path)
+            .map_err(|e| anyhow::anyhow!("Failed to load tokenizer: {}", e))?;
+        Ok(Self { tokenizer })
+    }
+
+    pub fn encode(&self, text: &str, add_special_tokens: bool) -> anyhow::Result<Vec<u32>> {
+        let encoding = self
+            .tokenizer
+            .encode(text, add_special_tokens)
+            .map_err(|e| anyhow::anyhow!("Failed to encode: {}", e))?;
+        Ok(encoding.get_ids().to_vec())
+    }
+
+    pub fn decode(&self, ids: &[u32], skip_special_tokens: bool) -> anyhow::Result<String> {
+        let text = self
+            .tokenizer
+            .decode(ids, skip_special_tokens)
+            .map_err(|e| anyhow::anyhow!("Failed to decode: {}", e))?;
+        Ok(text)
+    }
+
+    pub fn vocab_size(&self) -> usize {
+        self.tokenizer.get_vocab_size(true)
+    }
+
+    /// 应用 MiniCPM5 chat template
+    pub fn apply_chat_template(&self, user_msg: &str, enable_thinking: bool) -> anyhow::Result<Vec<u32>> {
+        let prompt = if enable_thinking {
+            format!(
+                "<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n",
+                user_msg
+            )
+        } else {
+            format!(
+                "<s><|im_start|>user\n{}<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n",
+                user_msg
+            )
+        };
+        self.encode(&prompt, false)
+    }
+}
+
+
+pub struct GenerationConfig {
+    pub max_new_tokens: Option<usize>,
+    pub temperature: f32,
+    pub top_p: f32,
+}
+
+impl Default for GenerationConfig {
+    fn default() -> Self {
+        Self {
+            max_new_tokens: None,
+            temperature: 1.0,
+            top_p: 1.0,
+        }
+    }
+}
+
+pub fn generate_with_cache<B: Backend>(
+    model: &LlamaForCausalLM<B>,
+    input_ids: &[u32],
+    config: &GenerationConfig,
+    eos_token_id: &EosTokenId,
+    device: &B::Device,
+) -> Vec<u32> {
+    let mut output_ids = input_ids.to_vec();
+    let mut cache: Option<LlamaKVCache<B>> = None;
+
+    // 第一步：输入完整 prompt，建立初始 cache
+    {
+        let input_ints: Vec<i64> = input_ids.iter().map(|&x| x as i64).collect();
+        let input_tensor =
+            Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
+                .unsqueeze::<2>();
+
+        let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
+        cache = Some(new_cache);
+
+        let next_token = sample_last(&logits, config.temperature, config.top_p);
+        output_ids.push(next_token);
+
+        if eos_token_id.contains(next_token) {
+            return output_ids;
+        }
+    }
+
+    // 后续步骤：每次只输入 1 个 token，使用 cache
+    let mut count = 1;
+    loop {
+        if let Some(max) = config.max_new_tokens {
+            if count >= max {
+                break;
+            }
+        }
+
+        let last_token = output_ids[output_ids.len() - 1];
+        let input_ints: Vec<i64> = vec![last_token as i64];
+        let input_tensor =
+            Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), device)
+                .unsqueeze::<2>();
+
+        let (logits, new_cache) = model.forward_with_cache(input_tensor, cache.as_ref());
+        cache = Some(new_cache);
+
+        let vocab_size = model.config.vocab_size;
+        let next_token_logits = logits
+            .clone()
+            .slice([0..1, 0..1, 0..vocab_size])
+            .reshape([vocab_size]);
+
+        let next_token = sample(&next_token_logits, config.temperature, config.top_p);
+        output_ids.push(next_token);
+
+        if eos_token_id.contains(next_token) {
+            break;
+        }
+
+        count += 1;
+    }
+
+    output_ids
+}
+
+fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
+    let shape = logits.shape();
+    let dims: [usize; 3] = shape.dims();
+    let seq_len = dims[1];
+    let vocab_size = dims[2];
+
+    let last_logits = logits
+        .clone()
+        .slice([0..1, seq_len - 1..seq_len, 0..vocab_size])
+        .reshape([vocab_size]);
+
+    sample(&last_logits, temperature, top_p)
+}
+
+fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p: f32) -> u32 {
+    let data = logits.clone().to_data();
+    let values: Vec<f32> = data.as_slice().unwrap().to_vec();
+
+    // temperature = 0 时退化为 greedy
+    if temperature <= 0.001 {
+        let mut max_idx = 0;
+        let mut max_val = f32::NEG_INFINITY;
+        for (i, &val) in values.iter().enumerate() {
+            if val > max_val {
+                max_val = val;
+                max_idx = i;
+            }
+        }
+        return max_idx as u32;
+    }
+
+    // 应用 temperature
+    let scaled: Vec<f32> = values.iter().map(|&v| v / temperature).collect();
+
+    // softmax
+    let max_val = scaled.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+    let exp_sum: f32 = scaled.iter().map(|&v| (v - max_val).exp()).sum();
+    let mut probs: Vec<f32> = scaled.iter().map(|&v| (v - max_val).exp() / exp_sum).collect();
+
+    // top_p 过滤
+    if top_p < 1.0 {
+        let mut indexed: Vec<(usize, f32)> = probs.iter().enumerate().map(|(i, &p)| (i, p)).collect();
+        indexed.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap());
+
+        let mut cumsum = 0.0;
+        let mut cutoff = indexed.len();
+        for (i, &(_, p)) in indexed.iter().enumerate() {
+            cumsum += p;
+            if cumsum > top_p {
+                cutoff = i + 1;
+                break;
+            }
+        }
+
+        let keep: std::collections::HashSet<usize> = indexed[..cutoff].iter().map(|(i, _)| *i).collect();
+        for (i, p) in probs.iter_mut().enumerate() {
+            if !keep.contains(&i) {
+                *p = 0.0;
+            }
+        }
+
+        let sum: f32 = probs.iter().sum();
+        if sum > 0.0 {
+            for p in probs.iter_mut() {
+                *p /= sum;
+            }
+        }
+    }
+
+    // 多项式采样
+    let mut rng = rand::thread_rng();
+    let r: f32 = rng.gen();
+    let mut cumsum = 0.0;
+    for (i, &p) in probs.iter().enumerate() {
+        cumsum += p;
+        if r < cumsum {
+            return i as u32;
+        }
+    }
+    probs.len() as u32 - 1
+}
+
+
+pub struct MiniCPM<B: Backend> {
+    model: LlamaForCausalLM<B>,
+    config: LlamaConfig,
+    tokenizer: TokenizerWrapper,
+    device: B::Device,
+}
+
+impl<B: Backend> MiniCPM<B> {
+    pub fn load(
+        model_path: &str,
+        config_path: &str,
+        tokenizer_path: &str,
+        device: &B::Device,
+    ) -> anyhow::Result<Self> {
+        let config = LlamaConfig::from_json(config_path)?;
+        let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
+
+        let model = LlamaForCausalLM::<B>::new(config.clone(), device);
+        let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
+        let model = model.load_file(model_path, &recorder, device)?;
+
+        Ok(Self {
+            model,
+            config,
+            tokenizer,
+            device: device.clone(),
+        })
+    }
+
+    pub fn generate(
+        &self,
+        prompt: &str,
+        think: bool,
+        config: &GenerationConfig,
+    ) -> anyhow::Result<String> {
+        let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
+        let output_ids = generate_with_cache(
+            &self.model,
+            &input_ids,
+            config,
+            &self.config.eos_token_id,
+            &self.device,
+        );
+        let new_ids = &output_ids[input_ids.len()..];
+        self.tokenizer.decode(new_ids, true)
+    }
+
+    pub fn config(&self) -> &LlamaConfig {
+        &self.config
+    }
+
+    pub fn tokenizer(&self) -> &TokenizerWrapper {
+        &self.tokenizer
+    }
+
+    pub fn inner_model(&self) -> &LlamaForCausalLM<B> {
+        &self.model
+    }
+}