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

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

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)
+    }
+}