chore(代码变更): 1. 清理无用代码 2. 更新依赖库 3. 优化项目结构 4. 修复小型bug 5. 增加注释以提高可读性

crates/minicpm-convert/Cargo.toml

@@ -9,4 +9,5 @@ minicpm-core = { path = "../minicpm-core", version = "0.1.2", registry = "gitea"
 burn = { version = "0.21", default-features = false, features = ["std"] }
 memmap2 = "0.9"
 anyhow = "1.0"
-serde_json = "1.0"
+serde_json = "1.0"
+serde = { version = "1.0", features = ["derive"] }

crates/minicpm-convert/src/lib.rs

@@ -2,7 +2,7 @@ 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::record::{FullPrecisionSettings, HalfPrecisionSettings, NamedMpkFileRecorder};
 use burn::tensor::backend::Backend;
 use burn::tensor::{Float, Shape, Tensor, TensorData};
 use memmap2::Mmap;
@@ -40,6 +40,39 @@ pub fn export_model<B: Backend>(
     println!("模型已成功导出到: {:?}", output_dir);
     Ok(())
 }
+
+/// 导出模型为半精度（f16），输出文件约为全精度的一半大小。
+pub fn export_model_half<B: Backend>(
+    safetensors_path: &Path,
+    config_path: &Path,
+    tokenizer_path: &Path,
+    output_dir: &Path,
+    device: &B::Device,
+) -> anyhow::Result<()> {
+    println!("开始转换 MiniCPM 模型为 Burn 格式（半精度 f16）...");
+
+    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::<HalfPrecisionSettings>::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,
@@ -256,3 +289,223 @@ fn f16_to_f32(f16: u16) -> f32 {
     let f32_bits = ((sign as u32) << 31) | (f32_exp << 23) | ((mant as u32) << 13);
     f32::from_bits(f32_bits)
 }
+
+// ==================== Q8 量化导出 ====================
+
+/// Q8 量化单个 f32 切片，返回 (scale, i8_data)
+fn quantize_q8(f32_data: &[f32]) -> (f32, Vec<i8>) {
+    let max_abs = f32_data
+        .iter()
+        .map(|&x| x.abs())
+        .fold(0.0f32, |a, b| a.max(b));
+    let scale = if max_abs > 0.0 { max_abs / 127.0 } else { 1.0 };
+    let q8: Vec<i8> = f32_data
+        .iter()
+        .map(|&x| (x / scale).round().clamp(-127.0, 127.0) as i8)
+        .collect();
+    (scale, q8)
+}
+
+/// Q8 导出元数据中的 tensor 条目
+#[derive(serde::Serialize, serde::Deserialize)]
+struct Q8TensorMeta {
+    shape: Vec<usize>,
+    scale: f32,
+    offset: u64,
+    numel: usize,
+}
+
+/// Q8 导出元数据
+#[derive(serde::Serialize, serde::Deserialize)]
+struct Q8Metadata {
+    tensors: std::collections::HashMap<String, Q8TensorMeta>,
+}
+
+/// 导出模型为 Q8 量化格式（INT8），输出 model.q8.json + model.q8.bin。
+/// 2D 权重（Linear/Embedding）量化为 INT8；1D norm 权重存为 F16。
+pub fn export_model_q8(
+    safetensors_path: &Path,
+    config_path: &Path,
+    tokenizer_path: &Path,
+    output_dir: &Path,
+) -> anyhow::Result<()> {
+    println!("开始转换 MiniCPM 模型为 Q8 量化格式...");
+
+    println!("解析 safetensors 文件: {:?}", safetensors_path);
+    let file = std::fs::File::open(safetensors_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 data_offset = 8 + header_len;
+
+    // 收集 tensor 信息
+    let mut tensor_infos: Vec<(String, Vec<usize>, String, usize, usize)> = Vec::new();
+    if let Some(obj) = header.as_object() {
+        for (name, info) in obj {
+            if name == "__metadata__" {
+                continue;
+            }
+            if let Some(info_obj) = info.as_object() {
+                let mut shape = Vec::new();
+                if let Some(s) = info_obj.get("shape") {
+                    if let Some(arr) = s.as_array() {
+                        shape = arr.iter().map(|v| v.as_u64().unwrap() as usize).collect();
+                    }
+                }
+                let mut dtype = String::new();
+                if let Some(d) = info_obj.get("dtype") {
+                    dtype = d.as_str().unwrap_or("").to_string();
+                }
+                let mut start = 0;
+                let mut end = 0;
+                if let Some(offsets) = info_obj.get("data_offsets") {
+                    if let Some(arr) = offsets.as_array() {
+                        start = arr[0].as_u64().unwrap() as usize;
+                        end = arr[1].as_u64().unwrap() as usize;
+                    }
+                }
+                tensor_infos.push((name.clone(), shape, dtype, start, end));
+            }
+        }
+    }
+
+    // 定义需要量化的 2D 权重名称（与 load_safetensors 中的列表对应）
+    let mut q8_tensors: Vec<String> = Vec::new();
+    for i in 0..32 {
+        let p = format!("model.layers.{i}");
+        q8_tensors.push(format!("{p}.self_attn.q_proj.weight"));
+        q8_tensors.push(format!("{p}.self_attn.k_proj.weight"));
+        q8_tensors.push(format!("{p}.self_attn.v_proj.weight"));
+        q8_tensors.push(format!("{p}.self_attn.o_proj.weight"));
+        q8_tensors.push(format!("{p}.mlp.gate_proj.weight"));
+        q8_tensors.push(format!("{p}.mlp.up_proj.weight"));
+        q8_tensors.push(format!("{p}.mlp.down_proj.weight"));
+    }
+    q8_tensors.push("model.embed_tokens.weight".to_string());
+    q8_tensors.push("lm_head.weight".to_string());
+
+    let q8_set: std::collections::HashSet<String> = q8_tensors.into_iter().collect();
+
+    std::fs::create_dir_all(output_dir)?;
+    let json_path = output_dir.join("model.q8.json");
+    let bin_path = output_dir.join("model.q8.bin");
+
+    let mut metadata: std::collections::HashMap<String, Q8TensorMeta> = std::collections::HashMap::new();
+    let mut bin_file = std::fs::File::create(&bin_path)?;
+    let mut current_offset: u64 = 0;
+
+    // 先处理 2D 权重：量化为 INT8
+    for (name, shape, dtype, start, end) in &tensor_infos {
+        if shape.len() != 2 {
+            continue;
+        }
+        if !q8_set.contains(name.as_str()) {
+            continue;
+        }
+
+        println!("量化 2D 权重: {name}  shape={shape:?}  dtype={dtype}");
+        let raw = &mmap[data_offset + *start..data_offset + *end];
+        let f32_data = convert_to_f32(raw, dtype);
+        let numel = f32_data.len();
+        let (scale, q8_data) = quantize_q8(&f32_data);
+
+        // 写入 INT8 数据
+        let bytes: &[u8] = unsafe {
+            std::slice::from_raw_parts(q8_data.as_ptr() as *const u8, q8_data.len())
+        };
+        use std::io::Write;
+        bin_file.write_all(bytes)?;
+
+        metadata.insert(
+            name.clone(),
+            Q8TensorMeta {
+                shape: shape.clone(),
+                scale,
+                offset: current_offset,
+                numel,
+            },
+        );
+        current_offset += q8_data.len() as u64;
+    }
+
+    // 再处理 1D norm 权重：存为 F16
+    for (name, shape, dtype, start, end) in &tensor_infos {
+        if shape.len() != 1 {
+            continue;
+        }
+        let raw = &mmap[data_offset + *start..data_offset + *end];
+        let f32_data = convert_to_f32(raw, dtype);
+        let numel = f32_data.len();
+
+        // 转 F16 存储
+        let mut f16_bytes = Vec::with_capacity(numel * 2);
+        for &v in &f32_data {
+            let f16_bits = f32_to_f16_bits(v);
+            f16_bytes.extend_from_slice(&f16_bits.to_le_bytes());
+        }
+
+        use std::io::Write;
+        bin_file.write_all(&f16_bytes)?;
+
+        // 用负数 scale 表示 F16 类型（加载时 scale < 0 表示 F16）
+        metadata.insert(
+            name.clone(),
+            Q8TensorMeta {
+                shape: shape.clone(),
+                scale: -1.0, // 标记为 F16
+                offset: current_offset,
+                numel,
+            },
+        );
+        current_offset += f16_bytes.len() as u64;
+    }
+
+    // 写 JSON 元数据
+    let meta = Q8Metadata { tensors: metadata };
+    let json_str = serde_json::to_string_pretty(&meta)?;
+    std::fs::write(&json_path, json_str)?;
+
+    println!("拷贝配置文件和 tokenizer...");
+    std::fs::copy(config_path, output_dir.join("config.json"))?;
+    std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
+
+    println!("Q8 量化模型已导出到: {:?}", output_dir);
+    println!("  - model.q8.json (元数据)",);
+    println!("  - model.q8.bin (INT8 + F16 权重)",);
+    Ok(())
+}
+
+/// f32 → f16 位表示（返回 u16）
+fn f32_to_f16_bits(v: f32) -> u16 {
+    let bits = v.to_bits();
+    let sign = bits >> 31;
+    let exp = (bits >> 23) & 0xFF;
+    let mant = bits & 0x7FFFFF;
+
+    let res: u32 = if exp == 0xFF {
+        // NaN / Inf
+        let f16_exp = 0x1F;
+        let f16_mant = if mant == 0 { 0 } else { 0x3FF };
+        (sign << 15) | (f16_exp << 10) | f16_mant
+    } else if exp < 103 {
+        // 下溢，返回 0
+        0
+    } else {
+        let new_exp = (exp as i32) - 127 + 15;
+        if new_exp >= 31 {
+            // 上溢
+            (sign << 15) | (0x1F << 10)
+        } else if new_exp <= 0 {
+            // 次正规数
+            let mant_new = (mant | 0x800000) >> (113 - new_exp);
+            (sign << 15) | (mant_new & 0x3FF)
+        } else {
+            let f16_mant = mant >> 13;
+            (sign << 15) | ((new_exp as u32) << 10) | f16_mant
+        }
+    };
+    res as u16
+}

crates/minicpm-inference/Cargo.toml

@@ -9,4 +9,7 @@ minicpm-core = { path = "../minicpm-core", version = "0.1.2", registry = "gitea"
 burn = { version = "0.21", default-features = false, features = ["std"] }
 tokenizers = "0.20"
 rand = "0.8"
-anyhow = "1.0"
+anyhow = "1.0"
+memmap2 = "0.9"
+serde = { version = "1.0", features = ["derive"] }
+serde_json = "1.0"

crates/minicpm-inference/src/config.rs

@@ -0,0 +1,16 @@
+#[derive(Debug, Clone)]
+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,
+        }
+    }
+}

crates/minicpm-inference/src/generator.rs

@@ -0,0 +1,139 @@
+use burn::tensor::backend::Backend;
+use burn::tensor::{Int, Tensor};
+
+use minicpm_core::config::EosTokenId;
+use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
+
+use crate::config::GenerationConfig;
+use crate::sampling::{sample, sample_last};
+
+/// 流式生成迭代器 —— 用户可以自行遍历处理每个 token
+///
+/// 用法示例：
+/// ```ignore
+/// let stream = TokenStream::new(model, input_ids, config, eos_token_id, device);
+/// for token in stream {
+///     // 自行处理每个 token
+///     println!("{}", token);
+/// }
+/// ```
+pub struct TokenStream<'a, B: Backend> {
+    model: &'a LlamaForCausalLM<B>,
+    config: &'a GenerationConfig,
+    eos_token_id: &'a EosTokenId,
+    device: &'a B::Device,
+    cache: Option<LlamaKVCache<B>>,
+    last_token: Option<u32>,
+    count: usize,
+    finished: bool,
+    first_step: bool,
+    input_ids: Vec<u32>,
+}
+
+impl<'a, B: Backend> TokenStream<'a, B> {
+    pub fn new(
+        model: &'a LlamaForCausalLM<B>,
+        input_ids: &[u32],
+        config: &'a GenerationConfig,
+        eos_token_id: &'a EosTokenId,
+        device: &'a B::Device,
+    ) -> Self {
+        Self {
+            model,
+            config,
+            eos_token_id,
+            device,
+            cache: None,
+            last_token: None,
+            count: 0,
+            finished: false,
+            first_step: true,
+            input_ids: input_ids.to_vec(),
+        }
+    }
+}
+
+impl<'a, B: Backend> Iterator for TokenStream<'a, B> {
+    type Item = u32;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        if self.finished {
+            return None;
+        }
+
+        if let Some(max) = self.config.max_new_tokens {
+            if self.count >= max {
+                self.finished = true;
+                return None;
+            }
+        }
+
+        let next_token = if self.first_step {
+            self.first_step = false;
+            self.step_first()
+        } else {
+            self.step_next()
+        };
+
+        match next_token {
+            Some(token) => {
+                self.count += 1;
+                self.last_token = Some(token);
+
+                if self.eos_token_id.contains(token) {
+                    self.finished = true;
+                }
+
+                Some(token)
+            }
+            None => {
+                self.finished = true;
+                None
+            }
+        }
+    }
+}
+
+impl<'a, B: Backend> TokenStream<'a, B> {
+    fn step_first(&mut self) -> Option<u32> {
+        let input_ints: Vec<i64> = self.input_ids.iter().map(|&x| x as i64).collect();
+        let input_tensor =
+            Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), self.device)
+                .unsqueeze::<2>();
+
+        let (logits, new_cache) = self.model.forward_with_cache(input_tensor, self.cache.as_ref());
+        self.cache = Some(new_cache);
+
+        Some(sample_last(&logits, self.config.temperature, self.config.top_p))
+    }
+
+    fn step_next(&mut self) -> Option<u32> {
+        let last_token = self.last_token?;
+        let input_ints: Vec<i64> = vec![last_token as i64];
+        let input_tensor =
+            Tensor::<B, 1, Int>::from_ints(input_ints.as_slice(), self.device)
+                .unsqueeze::<2>();
+
+        let (logits, new_cache) = self.model.forward_with_cache(input_tensor, self.cache.as_ref());
+        self.cache = Some(new_cache);
+
+        let vocab_size = self.model.config.vocab_size;
+        let next_token_logits = logits
+            .slice([0..1, 0..1, 0..vocab_size])
+            .reshape([vocab_size]);
+
+        Some(sample(&next_token_logits, self.config.temperature, self.config.top_p))
+    }
+}
+
+/// 非流式生成 —— 返回所有生成的 token id
+pub fn generate_tokens<B: Backend>(
+    model: &LlamaForCausalLM<B>,
+    input_ids: &[u32],
+    config: &GenerationConfig,
+    eos_token_id: &EosTokenId,
+    device: &B::Device,
+) -> Vec<u32> {
+    let stream = TokenStream::new(model, input_ids, config, eos_token_id, device);
+    stream.collect()
+}

crates/minicpm-inference/src/lib.rs

@@ -1,305 +1,18 @@
+pub mod config;
+pub mod generator;
+pub mod loader;
+pub mod sampling;
+pub mod tokenizer;
+
+pub use config::GenerationConfig;
+pub use generator::TokenStream;
 pub use minicpm_core::config::{EosTokenId, LlamaConfig};
 pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
+pub use tokenizer::TokenizerWrapper;
 
-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
-}
-
-/// 流式生成：每生成一个新 token，立即调用 `on_token` 回调输出该 token 的解码文本。
-/// 返回完整的 output token IDs。
-pub fn generate_stream<B: Backend>(
-    model: &LlamaForCausalLM<B>,
-    tokenizer: &TokenizerWrapper,
-    input_ids: &[u32],
-    config: &GenerationConfig,
-    eos_token_id: &EosTokenId,
-    device: &B::Device,
-    mut on_token: impl FnMut(&str),
-) -> Vec<u32> {
-    let mut output_ids = input_ids.to_vec();
-    let mut cache: Option<LlamaKVCache<B>> = None;
-
-    // 第一步：完整 prompt 输入
-    {
-        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);
-
-        // 流式输出第一个 token
-        if let Ok(text) = tokenizer.decode(&[next_token], true) {
-            on_token(&text);
-        }
-
-        if eos_token_id.contains(next_token) {
-            return output_ids;
-        }
-    }
-
-    // 后续步骤：每步生成一个 token
-    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);
-
-        // 流式输出新 token
-        if let Ok(text) = tokenizer.decode(&[next_token], true) {
-            on_token(&text);
-        }
-
-        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
-}
 
+use loader::{load_model, load_model_q8_from_dir};
 
 pub struct MiniCPM<B: Backend> {
     model: LlamaForCausalLM<B>,
@@ -317,10 +30,25 @@ impl<B: Backend> MiniCPM<B> {
     ) -> anyhow::Result<Self> {
         let config = LlamaConfig::from_json(config_path)?;
         let tokenizer = TokenizerWrapper::from_file(tokenizer_path)?;
+        let model = load_model::<B>(model_path, &config, device)?;
 
-        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 load_q8(
+        model_dir: &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 = load_model_q8_from_dir::<B>(model_dir, &config, device)?;
 
         Ok(Self {
             model,
@@ -337,7 +65,7 @@ impl<B: Backend> MiniCPM<B> {
         config: &GenerationConfig,
     ) -> anyhow::Result<String> {
         let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
-        let output_ids = generate_with_cache(
+        let output_ids = generator::generate_tokens(
             &self.model,
             &input_ids,
             config,
@@ -348,27 +76,26 @@ impl<B: Backend> MiniCPM<B> {
         self.tokenizer.decode(new_ids, true)
     }
 
-    /// 流式生成：每生成一个 token 立即调用 `on_token` 回调，
-    /// 参数为该 token 解码后的文本。
-    pub fn generate_stream(
-        &self,
+    pub fn generate_stream<'a>(
+        &'a self,
         prompt: &str,
         think: bool,
-        config: &GenerationConfig,
-        on_token: impl FnMut(&str),
-    ) -> anyhow::Result<String> {
+        config: &'a GenerationConfig,
+    ) -> anyhow::Result<TextStream<'a, B>> {
         let input_ids = self.tokenizer.apply_chat_template(prompt, think)?;
-        let output_ids = generate_stream(
+        let token_stream = TokenStream::new(
             &self.model,
-            &self.tokenizer,
             &input_ids,
             config,
             &self.config.eos_token_id,
             &self.device,
-            on_token,
         );
-        let new_ids = &output_ids[input_ids.len()..];
-        self.tokenizer.decode(new_ids, true)
+
+        Ok(TextStream {
+            token_stream,
+            tokenizer: &self.tokenizer,
+            buffer: Vec::new(),
+        })
     }
 
     pub fn config(&self) -> &LlamaConfig {
@@ -383,3 +110,41 @@ impl<B: Backend> MiniCPM<B> {
         &self.model
     }
 }
+
+/// 文本流式迭代器 —— 解码 token 并返回文本片段
+///
+/// 用法示例：
+/// ```ignore
+/// let mut stream = model.generate_stream("你好", false, &config)?;
+/// for text in stream {
+///     print!("{}", text);
+/// }
+/// ```
+pub struct TextStream<'a, B: Backend> {
+    token_stream: TokenStream<'a, B>,
+    tokenizer: &'a TokenizerWrapper,
+    buffer: Vec<u32>,
+}
+
+impl<'a, B: Backend> Iterator for TextStream<'a, B> {
+    type Item = String;
+
+    fn next(&mut self) -> Option<Self::Item> {
+        match self.token_stream.next() {
+            Some(token) => {
+                self.buffer.push(token);
+                match self.tokenizer.decode(&[token], true) {
+                    Ok(text) => Some(text),
+                    Err(_) => Some(String::new()),
+                }
+            }
+            None => None,
+        }
+    }
+}
+
+impl<'a, B: Backend> TextStream<'a, B> {
+    pub fn all_tokens(&self) -> &[u32] {
+        &self.buffer
+    }
+}

crates/minicpm-inference/src/loader/mod.rs

@@ -0,0 +1,29 @@
+pub mod q8;
+
+use burn::module::Module;
+use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
+use burn::tensor::backend::Backend;
+
+use minicpm_core::config::LlamaConfig;
+use minicpm_core::model::LlamaForCausalLM;
+
+use self::q8::load_model_q8;
+
+pub fn load_model<B: Backend>(
+    model_path: &str,
+    config: &LlamaConfig,
+    device: &B::Device,
+) -> anyhow::Result<LlamaForCausalLM<B>> {
+    let model = LlamaForCausalLM::<B>::new(config.clone(), device);
+    let recorder = NamedMpkFileRecorder::<FullPrecisionSettings>::new();
+    let model = model.load_file(model_path, &recorder, device)?;
+    Ok(model)
+}
+
+pub fn load_model_q8_from_dir<B: Backend>(
+    model_dir: &str,
+    config: &LlamaConfig,
+    device: &B::Device,
+) -> anyhow::Result<LlamaForCausalLM<B>> {
+    load_model_q8(model_dir, config, device)
+}

crates/minicpm-inference/src/loader/q8.rs

@@ -0,0 +1,201 @@
+use burn::module::{Module, Param};
+use burn::nn::{EmbeddingRecord, LinearRecord};
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Shape, Tensor, TensorData};
+use memmap2::Mmap;
+use std::collections::HashMap;
+
+use minicpm_core::config::LlamaConfig;
+use minicpm_core::model::LlamaForCausalLM;
+
+#[derive(serde::Deserialize)]
+struct Q8TensorMeta {
+    shape: Vec<usize>,
+    scale: f32,
+    offset: u64,
+    numel: usize,
+}
+
+#[derive(serde::Deserialize)]
+struct Q8Metadata {
+    tensors: HashMap<String, Q8TensorMeta>,
+}
+
+pub fn load_model_q8<B: Backend>(
+    model_dir: &str,
+    config: &LlamaConfig,
+    device: &B::Device,
+) -> anyhow::Result<LlamaForCausalLM<B>> {
+    let json_path = std::path::Path::new(model_dir).join("model.q8.json");
+    let bin_path = std::path::Path::new(model_dir).join("model.q8.bin");
+
+    let json_str = std::fs::read_to_string(&json_path)?;
+    let metadata: Q8Metadata = serde_json::from_str(&json_str)?;
+
+    let file = std::fs::File::open(&bin_path)?;
+    let mmap = unsafe { Mmap::map(&file)? };
+
+    let mut model = LlamaForCausalLM::<B>::new(config.clone(), device);
+
+    if let Some(meta) = metadata.tensors.get("model.embed_tokens.weight") {
+        let w = if meta.scale >= 0.0 {
+            load_q8_tensor_2d::<B>(&mmap, meta, device)
+        } else {
+            load_f16_tensor_2d::<B>(&mmap, meta, device)
+        };
+        let record = EmbeddingRecord { weight: Param::from_tensor(w) };
+        model.model.embed_tokens = model.model.embed_tokens.clone().load_record(record);
+    }
+
+    for i in 0..config.num_hidden_layers {
+        let p = format!("model.layers.{i}");
+
+        if let Some(meta) = metadata.tensors.get(&format!("{p}.self_attn.q_proj.weight")) {
+            let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
+            let record = LinearRecord { weight: Param::from_tensor(w), 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(meta) = metadata.tensors.get(&format!("{p}.self_attn.k_proj.weight")) {
+            let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
+            let record = LinearRecord { weight: Param::from_tensor(w), 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(meta) = metadata.tensors.get(&format!("{p}.self_attn.v_proj.weight")) {
+            let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
+            let record = LinearRecord { weight: Param::from_tensor(w), 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(meta) = metadata.tensors.get(&format!("{p}.self_attn.o_proj.weight")) {
+            let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
+            let record = LinearRecord { weight: Param::from_tensor(w), 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(meta) = metadata.tensors.get(&format!("{p}.mlp.gate_proj.weight")) {
+            let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
+            let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
+            model.model.layers[i].mlp.gate_proj = model.model.layers[i].mlp.gate_proj.clone().load_record(record);
+        }
+        if let Some(meta) = metadata.tensors.get(&format!("{p}.mlp.up_proj.weight")) {
+            let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
+            let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
+            model.model.layers[i].mlp.up_proj = model.model.layers[i].mlp.up_proj.clone().load_record(record);
+        }
+        if let Some(meta) = metadata.tensors.get(&format!("{p}.mlp.down_proj.weight")) {
+            let w = load_q8_tensor_2d::<B>(&mmap, meta, device).transpose();
+            let record = LinearRecord { weight: Param::from_tensor(w), bias: None };
+            model.model.layers[i].mlp.down_proj = model.model.layers[i].mlp.down_proj.clone().load_record(record);
+        }
+
+        if let Some(meta) = metadata.tensors.get(&format!("{p}.input_layernorm.weight")) {
+            let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
+            model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
+        }
+        if let Some(meta) = metadata.tensors.get(&format!("{p}.post_attention_layernorm.weight")) {
+            let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
+            model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
+        }
+    }
+
+    if let Some(meta) = metadata.tensors.get("model.norm.weight") {
+        let w = load_f16_tensor_1d::<B>(&mmap, meta, device);
+        model.model.norm.weight = Param::from_tensor(w);
+    }
+
+    if let Some(meta) = metadata.tensors.get("lm_head.weight") {
+        let w = if meta.scale >= 0.0 {
+            load_q8_tensor_2d::<B>(&mmap, meta, device)
+        } else {
+            load_f16_tensor_2d::<B>(&mmap, meta, device)
+        };
+        let w_t = w.transpose();
+        let record = LinearRecord { weight: Param::from_tensor(w_t), bias: None };
+        model.lm_head = model.lm_head.clone().load_record(record);
+    } else if config.tie_word_embeddings {
+        let embed_weight = model.model.embed_tokens.weight.val();
+        let w_t = embed_weight.transpose();
+        let record = LinearRecord { weight: Param::from_tensor(w_t), bias: None };
+        model.lm_head = model.lm_head.clone().load_record(record);
+    }
+
+    Ok(model)
+}
+
+fn load_q8_tensor_2d<B: Backend>(
+    mmap: &Mmap,
+    meta: &Q8TensorMeta,
+    device: &B::Device,
+) -> Tensor<B, 2, Float> {
+    let bytes = &mmap[meta.offset as usize..meta.offset as usize + meta.numel];
+    let q8: &[i8] = unsafe { std::slice::from_raw_parts(bytes.as_ptr() as *const i8, meta.numel) };
+    let f32_data: Vec<f32> = q8.iter().map(|&q| (q as f32) * meta.scale).collect();
+    let shape: [usize; 2] = [meta.shape[0], meta.shape[1]];
+    Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
+}
+
+fn load_f16_tensor_1d<B: Backend>(
+    mmap: &Mmap,
+    meta: &Q8TensorMeta,
+    device: &B::Device,
+) -> Tensor<B, 1, Float> {
+    let start = meta.offset as usize;
+    let bytes = &mmap[start..start + meta.numel * 2];
+    let f32_data: Vec<f32> = bytes
+        .chunks_exact(2)
+        .map(|c| {
+            let bits = u16::from_le_bytes([c[0], c[1]]);
+            f16_to_f32(bits)
+        })
+        .collect();
+    let shape: [usize; 1] = [meta.shape[0]];
+    Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
+}
+
+fn load_f16_tensor_2d<B: Backend>(
+    mmap: &Mmap,
+    meta: &Q8TensorMeta,
+    device: &B::Device,
+) -> Tensor<B, 2, Float> {
+    let start = meta.offset as usize;
+    let bytes = &mmap[start..start + meta.numel * 2];
+    let f32_data: Vec<f32> = bytes
+        .chunks_exact(2)
+        .map(|c| {
+            let bits = u16::from_le_bytes([c[0], c[1]]);
+            f16_to_f32(bits)
+        })
+        .collect();
+    let shape: [usize; 2] = [meta.shape[0], meta.shape[1]];
+    Tensor::from_data(TensorData::new(f32_data, Shape::new(shape)), device)
+}
+
+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-inference/src/sampling.rs

@@ -0,0 +1,80 @@
+use burn::tensor::backend::Backend;
+use burn::tensor::{Float, Tensor};
+use rand::Rng;
+
+pub 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)
+}
+
+pub 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();
+
+    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;
+    }
+
+    let scaled: Vec<f32> = values.iter().map(|&v| v / temperature).collect();
+
+    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();
+
+    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
+}

crates/minicpm-inference/src/tokenizer.rs

@@ -0,0 +1,48 @@
+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!("加载 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!("编码失败: {}", 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!("解码失败: {}", e))?;
+        Ok(text)
+    }
+
+    pub fn vocab_size(&self) -> usize {
+        self.tokenizer.get_vocab_size(true)
+    }
+
+    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)
+    }
+}