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

Cargo.toml

@@ -3,8 +3,11 @@ members = [
   "crates/minicpm-core",
   "crates/minicpm-convert",
   "crates/minicpm-inference",
-  "examples/convert",
-  "examples/flex-backend",
   "examples/wgpu-backend",
 ]
 resolver = "2"
+
+[profile.release]
+opt-level = 3
+lto = true
+codegen-units = 1

MODEL_GUIDE.md

@@ -0,0 +1,155 @@
+# 模型格式与使用指南
+
+## 模型格式说明
+
+项目支持三种模型格式，适用于不同场景：
+
+| 格式 | 精度 | 文件大小 | 适用后端 | 速度 |
+|------|------|----------|----------|------|
+| full | f32 全精度 | ~2.5GB | Flex / Wgpu | 基准 |
+| half | f16 半精度 | ~1.2GB | Flex / Wgpu | 较快 |
+| q8 | INT8 量化 | ~650MB | Flex (推荐) | 最快 |
+
+### 1. 全精度模型 (full)
+- 目录：`model/`
+- 文件：`model.mpk`
+- 精度：FP32
+- 适用场景：精度优先、GPU 推理
+
+### 2. 半精度模型 (half)
+- 目录：`model_half/`
+- 文件：`model.mpk`
+- 精度：FP16
+- 适用场景：GPU 推理、显存有限
+
+### 3. INT8 量化模型 (q8)
+- 目录：`model_q8/`
+- 文件：`model.q8.bin` + `model.q8.json`
+- 精度：INT8 + 缩放因子
+- 适用场景：CPU 推理、内存有限
+- 特性：SIMD 指令集加速（AVX2/SSE）
+
+## 使用方法
+
+### 命令行参数
+
+所有示例都支持以下命令行参数：
+
+```bash
+--full   # 使用全精度模型
+--half   # 使用半精度模型
+--q8     # 使用 INT8 量化模型
+```
+
+**重要**：参数前需要加 `--` 与 cargo 分隔：
+
+```bash
+# 正确
+cargo run --release -p flex-backend -- --q8
+cargo run --release -p wgpu-backend -- --half
+
+# 错误（参数会被 cargo 吃掉）
+cargo run --release -p flex-backend --q8
+```
+
+### Flex 后端（CPU）
+
+```bash
+# 全精度
+cargo run --release -p flex-backend -- --full
+
+# 半精度
+cargo run --release -p flex-backend -- --half
+
+# INT8 量化（推荐，最快）
+cargo run --release -p flex-backend -- --q8
+```
+
+**最佳性能编译**：
+```bash
+RUSTFLAGS="-C target-cpu=native" cargo build --release -p flex-backend
+cargo run --release -p flex-backend -- --q8
+```
+
+### Wgpu 后端（GPU）
+
+```bash
+# 全精度（推荐）
+cargo run --release -p wgpu-backend -- --full
+
+# 半精度
+cargo run --release -p wgpu-backend -- --half
+
+# INT8 量化（不推荐，性能较差）
+cargo run --release -p wgpu-backend -- --q8
+```
+
+**注意**：
+- WGPU 后端下，q8 量化层使用 CPU 实现，会频繁在 CPU-GPU 间传输数据，性能可能不如全精度
+- GPU 推理建议使用 `--full` 或 `--half`
+
+## 模型转换
+
+使用 convert 示例转换模型：
+
+```bash
+# 全精度
+cargo run --release -p convert -- --full
+
+# 半精度
+cargo run --release -p convert -- --half
+
+# INT8 量化
+cargo run --release -p convert -- --q8
+
+# 全部转换
+cargo run --release -p convert -- --all
+```
+
+源模型目录：`MiniCPM5-1B/`（需要包含 safetensors 格式模型）
+
+## 性能对比（参考）
+
+以 MiniCPM5-1B 模型为例：
+
+| 后端 | 格式 | 内存占用 | 推理速度 | 推荐指数 |
+|------|------|----------|----------|----------|
+| Flex CPU | q8 | ~1GB | ⭐⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ |
+| Flex CPU | half | ~1.5GB | ⭐⭐⭐ | ⭐⭐ |
+| Flex CPU | full | ~3GB | ⭐⭐ | ⭐ |
+| Wgpu GPU | full | ~2.5GB | ⭐⭐⭐⭐ | ⭐⭐⭐⭐ |
+| Wgpu GPU | half | ~1.5GB | ⭐⭐⭐⭐ | ⭐⭐⭐⭐⭐ |
+| Wgpu GPU | q8 | ~1GB | ⭐ | ⭐ |
+
+> 注：实际性能取决于 CPU/GPU 型号、内存带宽等因素
+
+## 启动时的输出说明
+
+程序启动时会打印配置信息，确认使用的模型格式：
+
+```
+========================================
+后端: Flex (纯 Rust CPU)
+模式: INT8 量化 (SIMD 加速)
+模型目录: model_q8
+配置文件: MiniCPM5-1B/config.json
+  特性: INT8 量化 + SIMD 加速，推荐使用
+  编译建议: RUSTFLAGS="-C target-cpu=native" cargo build --release
+========================================
+```
+
+如果模式和模型目录与预期不符，请检查命令行参数是否正确传递。
+
+## 常见问题
+
+### Q: 加了 --q8 参数但感觉没变化？
+A: 检查是否加了 `--` 分隔符。正确用法：`cargo run --release -p flex-backend -- --q8`
+
+### Q: WGPU 上 q8 反而更慢？
+A: 正常现象。q8 量化是 CPU 优化的，GPU 上建议用 full 或 half。
+
+### Q: half 模型速度和 full 差不多？
+A: Flex CPU 后端上 f16 运算会转成 f32，速度差异不大。WGPU GPU 后端上显存占用会明显减少。
+
+### Q: 怎么确认用的是哪个模型？
+A: 看启动时的打印信息，"模式"和"模型目录"字段会显示当前配置。

crates/minicpm-convert/src/format/full.rs

@@ -0,0 +1,31 @@
+use burn::module::Module;
+use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
+use burn::tensor::backend::Backend;
+use std::path::Path;
+
+use crate::loader::{load_into_model, TensorStore};
+use minicpm_core::config::LlamaConfig;
+use minicpm_core::model::LlamaForCausalLM;
+
+/// 导出全精度模型（f32）
+pub fn export<B: Backend>(
+    store: &TensorStore,
+    config: &LlamaConfig,
+    output_dir: &Path,
+    device: &B::Device,
+) -> anyhow::Result<()> {
+    println!("创建模型结构...");
+    let model = LlamaForCausalLM::<B>::new(config.clone(), device);
+
+    println!("加载 safetensors 权重...");
+    let model = load_into_model(model, store, 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!("全精度模型已导出到: {:?}", output_dir);
+    Ok(())
+}

crates/minicpm-convert/src/format/mod.rs

@@ -0,0 +1,44 @@
+pub mod full;
+
+/// 导出格式
+#[derive(Debug, Clone, Copy, PartialEq)]
+pub enum Format {
+    /// 全精度 f32
+    Full,
+}
+
+impl Format {
+    pub fn name(&self) -> &'static str {
+        match self {
+            Format::Full => "全精度",
+        }
+    }
+
+    pub fn dir_name(&self) -> &'static str {
+        match self {
+            Format::Full => "model",
+        }
+    }
+}
+
+/// 导出任务配置
+pub struct ExportTask {
+    pub format: Format,
+    pub output_dir: String,
+}
+
+impl ExportTask {
+    pub fn new(format: Format) -> Self {
+        Self {
+            output_dir: format.dir_name().to_string(),
+            format,
+        }
+    }
+
+    pub fn with_dir(format: Format, output_dir: impl Into<String>) -> Self {
+        Self {
+            output_dir: output_dir.into(),
+            format,
+        }
+    }
+}

crates/minicpm-convert/src/lib.rs

@@ -1,511 +1,46 @@
-use minicpm_core::config::LlamaConfig;
-use minicpm_core::model::LlamaForCausalLM;
-use burn::module::{Module, Param};
-use burn::nn::{EmbeddingRecord, LinearRecord};
-use burn::record::{FullPrecisionSettings, HalfPrecisionSettings, NamedMpkFileRecorder};
+pub mod format;
+pub mod loader;
+pub mod utils;
+
+pub use format::{ExportTask, Format};
+pub use loader::TensorStore;
+
 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>(
+use minicpm_core::config::LlamaConfig;
+
+/// 执行导出任务
+pub fn run_export<B: Backend>(
     safetensors_path: &Path,
     config_path: &Path,
     tokenizer_path: &Path,
-    output_dir: &Path,
+    tasks: &[ExportTask],
     device: &B::Device,
 ) -> anyhow::Result<()> {
-    println!("开始转换 MiniCPM 模型为 Burn 格式...");
+    println!("开始转换 MiniCPM 模型...");
+    println!("源文件: {:?}", safetensors_path);
 
-    println!("加载配置文件: {:?}", config_path);
+    println!("加载配置文件...");
     let config = LlamaConfig::from_json(config_path.to_str().unwrap())?;
 
-    println!("创建模型结构...");
-    let model = LlamaForCausalLM::<B>::new(config, device);
+    println!("解析 safetensors 文件...");
+    let store = TensorStore::from_file(safetensors_path)?;
 
-    println!("加载 safetensors 权重...");
-    let model = load_safetensors(model, safetensors_path, device)?;
+    for task in tasks {
+        println!("\n--- 导出 {} ---", task.format.name());
+        let output_dir = Path::new(&task.output_dir);
 
-    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)?;
+        match task.format {
+            Format::Full => {
+                format::full::export::<B>(&store, &config, output_dir, device)?;
+            }
+        }
 
-    println!("拷贝配置文件和 tokenizer...");
-    std::fs::copy(config_path, output_dir.join("config.json"))?;
-    std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
+        std::fs::copy(config_path, output_dir.join("config.json"))?;
+        std::fs::copy(tokenizer_path, output_dir.join("tokenizer.json"))?;
+    }
 
-    println!("模型已成功导出到: {:?}", output_dir);
+    println!("\n所有导出任务完成！");
     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,
-    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)
-}
-
-// ==================== 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-convert/src/loader/mod.rs

@@ -0,0 +1,3 @@
+pub mod safetensors;
+
+pub use safetensors::{load_into_model, TensorStore};

crates/minicpm-convert/src/loader/safetensors.rs

@@ -0,0 +1,185 @@
+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 std::path::Path;
+
+use crate::utils::convert_to_f32;
+use minicpm_core::model::LlamaForCausalLM;
+
+/// 从 safetensors 文件解析出的张量数据
+pub struct TensorStore {
+    tensors: HashMap<String, Vec<u8>>,
+    shapes: HashMap<String, Vec<usize>>,
+    dtypes: HashMap<String, String>,
+}
+
+impl TensorStore {
+    /// 从 safetensors 文件加载所有张量
+    pub fn from_file(path: &Path) -> anyhow::Result<Self> {
+        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());
+                        }
+                    }
+                }
+            }
+        }
+
+        Ok(Self {
+            tensors,
+            shapes,
+            dtypes,
+        })
+    }
+
+    /// 加载 2D 张量
+    pub fn load_2d<B: Backend>(&self, name: &str, device: &B::Device) -> Option<Tensor<B, 2, Float>> {
+        let data = self.tensors.get(name)?;
+        let shape = self.shapes.get(name)?;
+        let dtype = self.dtypes.get(name)?;
+
+        assert_eq!(shape.len(), 2, "期望 2D 张量: {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))
+    }
+
+    /// 加载 1D 张量
+    pub fn load_1d<B: Backend>(&self, name: &str, device: &B::Device) -> Option<Tensor<B, 1, Float>> {
+        let data = self.tensors.get(name)?;
+        let shape = self.shapes.get(name)?;
+        let dtype = self.dtypes.get(name)?;
+
+        assert_eq!(shape.len(), 1, "期望 1D 张量: {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))
+    }
+
+    /// 获取原始字节数据
+    pub fn raw(&self, name: &str) -> Option<&[u8]> {
+        self.tensors.get(name).map(|v| v.as_slice())
+    }
+
+    pub fn shape(&self, name: &str) -> Option<&Vec<usize>> {
+        self.shapes.get(name)
+    }
+
+    pub fn dtype(&self, name: &str) -> Option<&String> {
+        self.dtypes.get(name)
+    }
+
+    pub fn names(&self) -> impl Iterator<Item = &String> {
+        self.tensors.keys()
+    }
+}
+
+/// 将 safetensors 权重加载到模型中
+pub fn load_into_model<B: Backend>(
+    model: LlamaForCausalLM<B>,
+    store: &TensorStore,
+    device: &B::Device,
+) -> LlamaForCausalLM<B> {
+    let mut model = model;
+
+    if let Some(w) = store.load_2d("model.embed_tokens.weight", 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..model.config.num_hidden_layers {
+        let prefix = format!("model.layers.{i}");
+
+        load_linear(&mut model, store, device, i, &format!("{prefix}.self_attn.q_proj.weight"), |m| &mut m.model.layers[i].self_attn.q_proj);
+        load_linear(&mut model, store, device, i, &format!("{prefix}.self_attn.k_proj.weight"), |m| &mut m.model.layers[i].self_attn.k_proj);
+        load_linear(&mut model, store, device, i, &format!("{prefix}.self_attn.v_proj.weight"), |m| &mut m.model.layers[i].self_attn.v_proj);
+        load_linear(&mut model, store, device, i, &format!("{prefix}.self_attn.o_proj.weight"), |m| &mut m.model.layers[i].self_attn.o_proj);
+        load_linear(&mut model, store, device, i, &format!("{prefix}.mlp.gate_proj.weight"), |m| &mut m.model.layers[i].mlp.gate_proj);
+        load_linear(&mut model, store, device, i, &format!("{prefix}.mlp.up_proj.weight"), |m| &mut m.model.layers[i].mlp.up_proj);
+        load_linear(&mut model, store, device, i, &format!("{prefix}.mlp.down_proj.weight"), |m| &mut m.model.layers[i].mlp.down_proj);
+
+        if let Some(w) = store.load_1d(&format!("{prefix}.input_layernorm.weight"), device) {
+            model.model.layers[i].input_layernorm.weight = Param::from_tensor(w);
+        }
+        if let Some(w) = store.load_1d(&format!("{prefix}.post_attention_layernorm.weight"), device) {
+            model.model.layers[i].post_attention_layernorm.weight = Param::from_tensor(w);
+        }
+    }
+
+    if let Some(w) = store.load_1d("model.norm.weight", device) {
+        model.model.norm.weight = Param::from_tensor(w);
+    }
+
+    if let Some(w) = store.load_2d("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);
+    }
+
+    model
+}
+
+fn load_linear<B: Backend, F>(
+    model: &mut LlamaForCausalLM<B>,
+    store: &TensorStore,
+    device: &B::Device,
+    _layer_idx: usize,
+    name: &str,
+    get_field: F,
+) where
+    F: FnOnce(&mut LlamaForCausalLM<B>) -> &mut burn::nn::Linear<B>,
+{
+    if let Some(w) = store.load_2d(name, device) {
+        let record = LinearRecord { weight: Param::from_tensor(w.transpose()), bias: None };
+        let field = get_field(model);
+        *field = field.clone().load_record(record);
+    }
+}

crates/minicpm-convert/src/utils.rs

@@ -0,0 +1,84 @@
+/// BF16 → F32
+pub fn bf16_to_f32(bf16: u16) -> f32 {
+    let f32_bits = (bf16 as u32) << 16;
+    f32::from_bits(f32_bits)
+}
+
+/// F16 → F32
+pub 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)
+}
+
+/// F32 → F16 位表示
+pub 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 {
+        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
+    } 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
+}
+
+/// 将原始字节按 dtype 转换为 f32 向量
+pub fn convert_to_f32(data: &[u8], dtype: &str) -> Vec<f32> {
+    match dtype {
+        "BF16" => data
+            .chunks_exact(2)
+            .map(|chunk| bf16_to_f32(u16::from_le_bytes(chunk.try_into().unwrap())))
+            .collect(),
+        "F16" => data
+            .chunks_exact(2)
+            .map(|chunk| f16_to_f32(u16::from_le_bytes(chunk.try_into().unwrap())))
+            .collect(),
+        "F32" => data
+            .chunks_exact(4)
+            .map(|chunk| f32::from_le_bytes(chunk.try_into().unwrap()))
+            .collect(),
+        _ => panic!("不支持的 dtype: {dtype}"),
+    }
+}

crates/minicpm-core/Cargo.toml

@@ -4,6 +4,10 @@ version = "0.1.2"
 edition = "2021"
 publish = ["gitea"]
 
+[features]
+default = ["avx2"]
+avx2 = []
+
 [dependencies]
 burn = { version = "0.21", default-features = false, features = ["std"] }
 serde = { version = "1.0", features = ["derive"] }

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

@@ -129,7 +129,6 @@ impl<B: Backend> Attention<B> {
         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);

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

@@ -4,15 +4,15 @@ use burn::tensor::{Float, Tensor};
 
 use super::attention::{Attention, KVCache};
 use super::ffn::FeedForward;
-use super::norm::RmsNorm;
+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>,
+    pub input_layernorm: RMSNorm<B>,
+    pub post_attention_layernorm: RMSNorm<B>,
 }
 
 impl<B: Backend> DecoderLayer<B> {
@@ -27,8 +27,8 @@ impl<B: Backend> DecoderLayer<B> {
     ) -> 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);
+        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,

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

@@ -6,7 +6,7 @@ use burn::tensor::{Float, Int, Tensor};
 use super::attention::KVCache;
 use crate::config::LlamaConfig;
 use super::decoder::DecoderLayer;
-use super::norm::RmsNorm;
+use super::norm::RMSNorm;
 use super::rope::RoPE;
 
 #[derive(Debug, Clone)]
@@ -26,7 +26,7 @@ impl<B: Backend> LlamaKVCache<B> {
 pub struct LlamaModel<B: Backend> {
     pub embed_tokens: Embedding<B>,
     pub layers: Vec<DecoderLayer<B>>,
-    pub norm: RmsNorm<B>,
+    pub norm: RMSNorm<B>,
     pub rope: RoPE,
 }
 
@@ -58,7 +58,7 @@ impl<B: Backend> LlamaForCausalLM<B> {
             ));
         }
 
-        let norm = RmsNorm::new(config.hidden_size, 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);
 

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

@@ -4,12 +4,12 @@ use burn::tensor::backend::Backend;
 use burn::tensor::{Float, Tensor};
 
 #[derive(Module, Debug)]
-pub struct RmsNorm<B: Backend> {
+pub struct RMSNorm<B: Backend> {
     pub weight: Param<Tensor<B, 1, Float>>,
     pub eps: f64,
 }
 
-impl<B: Backend> RmsNorm<B> {
+impl<B: Backend> RMSNorm<B> {
     pub fn new(dim: usize, eps: f64, device: &B::Device) -> Self {
         let weight = Initializer::Ones.init([dim], device);
         Self {

crates/minicpm-inference/Cargo.toml

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

crates/minicpm-inference/src/generator.rs

@@ -5,18 +5,8 @@ use minicpm_core::config::EosTokenId;
 use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
 
 use crate::config::GenerationConfig;
-use crate::sampling::{sample, sample_last};
+use crate::sampling::{CachedSampler, 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,
@@ -28,6 +18,7 @@ pub struct TokenStream<'a, B: Backend> {
     finished: bool,
     first_step: bool,
     input_ids: Vec<u32>,
+    sampler: CachedSampler,
 }
 
 impl<'a, B: Backend> TokenStream<'a, B> {
@@ -49,6 +40,7 @@ impl<'a, B: Backend> TokenStream<'a, B> {
             finished: false,
             first_step: true,
             input_ids: input_ids.to_vec(),
+            sampler: CachedSampler::new(),
         }
     }
 }
@@ -122,11 +114,12 @@ impl<'a, B: Backend> TokenStream<'a, B> {
             .slice([0..1, 0..1, 0..vocab_size])
             .reshape([vocab_size]);
 
-        Some(sample(&next_token_logits, self.config.temperature, self.config.top_p))
+        let data = next_token_logits.to_data();
+        let values: Vec<f32> = data.as_slice().unwrap().to_vec();
+        Some(self.sampler.sample(&values, self.config.temperature, self.config.top_p))
     }
 }
 
-/// 非流式生成 —— 返回所有生成的 token id
 pub fn generate_tokens<B: Backend>(
     model: &LlamaForCausalLM<B>,
     input_ids: &[u32],

crates/minicpm-inference/src/lib.rs

@@ -5,14 +5,14 @@ pub mod sampling;
 pub mod tokenizer;
 
 pub use config::GenerationConfig;
-pub use generator::TokenStream;
+pub use generator::{TokenStream};
 pub use minicpm_core::config::{EosTokenId, LlamaConfig};
 pub use minicpm_core::model::{LlamaForCausalLM, LlamaKVCache};
 pub use tokenizer::TokenizerWrapper;
 
 use burn::tensor::backend::Backend;
 
-use loader::{load_model, load_model_q8_from_dir};
+use loader::load_model;
 
 pub struct MiniCPM<B: Backend> {
     model: LlamaForCausalLM<B>,
@@ -40,24 +40,6 @@ impl<B: Backend> MiniCPM<B> {
         })
     }
 
-    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,
-            config,
-            tokenizer,
-            device: device.clone(),
-        })
-    }
-
     pub fn generate(
         &self,
         prompt: &str,
@@ -111,15 +93,6 @@ impl<B: Backend> MiniCPM<B> {
     }
 }
 
-/// 文本流式迭代器 —— 解码 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,

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

@@ -1,5 +1,3 @@
-pub mod q8;
-
 use burn::module::Module;
 use burn::record::{FullPrecisionSettings, NamedMpkFileRecorder};
 use burn::tensor::backend::Backend;
@@ -7,8 +5,6 @@ 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,
@@ -19,11 +15,3 @@ pub fn load_model<B: Backend>(
     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

@@ -1,201 +0,0 @@
-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

@@ -1,6 +1,91 @@
 use burn::tensor::backend::Backend;
 use burn::tensor::{Float, Tensor};
-use rand::Rng;
+use rand::{rngs::ThreadRng, Rng};
+
+struct Sampler {
+    rng: ThreadRng,
+}
+
+impl Sampler {
+    fn new() -> Self {
+        Self {
+            rng: rand::thread_rng(),
+        }
+    }
+
+    fn sample(&mut self, logits: &[f32], temperature: f32, top_p: f32) -> u32 {
+        if temperature <= 0.001 {
+            let mut max_idx = 0;
+            let mut max_val = f32::NEG_INFINITY;
+            for (i, &val) in logits.iter().enumerate() {
+                if val > max_val {
+                    max_val = val;
+                    max_idx = i;
+                }
+            }
+            return max_idx as u32;
+        }
+
+        let max_val = logits.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+        let exp_sum: f32 = logits.iter().map(|&v| (v - max_val).exp()).sum();
+        let mut probs: Vec<f32> = logits.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 r: f32 = self.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 CachedSampler {
+    sampler: std::cell::RefCell<Sampler>,
+}
+
+impl CachedSampler {
+    pub fn new() -> Self {
+        Self {
+            sampler: std::cell::RefCell::new(Sampler::new()),
+        }
+    }
+
+    pub fn sample(&self, logits: &[f32], temperature: f32, top_p: f32) -> u32 {
+        self.sampler.borrow_mut().sample(logits, temperature, top_p)
+    }
+}
 
 pub fn sample_last<B: Backend>(logits: &Tensor<B, 3, Float>, temperature: f32, top_p: f32) -> u32 {
     let shape = logits.shape();
@@ -32,11 +117,9 @@ pub fn sample<B: Backend>(logits: &Tensor<B, 1, Float>, temperature: f32, top_p:
         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();
+    let max_val = values.iter().cloned().fold(f32::NEG_INFINITY, f32::max);
+    let exp_sum: f32 = values.iter().map(|&v| (v - max_val).exp()).sum();
+    let mut probs: Vec<f32> = values.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();

examples/convert/Cargo.toml

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

examples/convert/src/main.rs

@@ -1,14 +0,0 @@
-use minicpm_convert::export_model_q8;
-use std::path::Path;
-
-fn main() -> anyhow::Result<()> {
-    let safetensors_path = Path::new("MiniCPM5-1B/model-00000-of-00001.safetensors");
-    let config_path = Path::new("MiniCPM5-1B/config.json");
-    let tokenizer_path = Path::new("MiniCPM5-1B/tokenizer.json");
-
-    // Q8 量化导出（INT8，文件约 1/4 大小）
-    export_model_q8(safetensors_path, config_path, tokenizer_path, Path::new("model_q8"))?;
-
-    println!("模型转换完成！");
-    Ok(())
-}

examples/flex-backend/Cargo.toml

@@ -1,9 +0,0 @@
-[package]
-name = "flex-backend"
-version = "0.1.0"
-edition = "2021"
-
-[dependencies]
-minicpm-inference = { path = "../../crates/minicpm-inference" }
-burn = { version = "0.21", default-features = false, features = ["std", "flex"] }
-anyhow = "1.0"

examples/flex-backend/src/main.rs

@@ -1,82 +0,0 @@
-// Flex 后端推理示例 —— 适用于 Tauri / WebAssembly 环境
-//
-// burn-flex 是纯 Rust CPU 后端，无 GPU 依赖，完美适配 Tauri 应用。
-// 支持 SIMD 加速（AVX2/NEON）和多线程（rayon）。
-//
-// 运行方式：cargo run --release -p flex-backend [--full|--half|--q8]
-
-use burn::backend::{flex::FlexDevice, Flex};
-use minicpm_inference::{GenerationConfig, MiniCPM};
-use std::io::Write;
-use std::time::Instant;
-
-fn main() -> anyhow::Result<()> {
-    let args: Vec<String> = std::env::args().collect();
-    let mode = args.get(1).map(|s| s.as_str()).unwrap_or("--full");
-
-    let (model_dir, config_path, tokenizer_path, is_q8) = match mode {
-        "--half" => (
-            "model_half",
-            "MiniCPM5-1B/config.json",
-            "MiniCPM5-1B/tokenizer.json",
-            false,
-        ),
-        "--q8" => (
-            "model_q8",
-            "MiniCPM5-1B/config.json",
-            "MiniCPM5-1B/tokenizer.json",
-            true,
-        ),
-        _ => (
-            "model",
-            "MiniCPM5-1B/config.json",
-            "MiniCPM5-1B/tokenizer.json",
-            false,
-        ),
-    };
-
-    println!("正在加载模型（后端：Flex，模式：{mode}）...");
-    println!("Flex 后端特性：纯 Rust、SIMD 加速、多线程");
-    let start = Instant::now();
-
-    let device = FlexDevice::default();
-    let model: MiniCPM<Flex> = if is_q8 {
-        MiniCPM::<Flex>::load_q8(model_dir, config_path, tokenizer_path, &device)?
-    } else {
-        let model_path = format!("{model_dir}/model");
-        MiniCPM::<Flex>::load(&model_path, config_path, tokenizer_path, &device)?
-    };
-
-    println!("模型加载完成，耗时: {:.2?}", start.elapsed());
-
-    let prompt = "需要处理生成一个河南的旅游计划，要求如下：
-1. 旅游计划为三天两夜的行程安排。
-2. 每天的行程安排包括景点、餐饮和住宿建议。
-3. 景点安排要考虑交通便利性和时间合理性。
-请生成详细的旅游计划，包含每天的行程安排、景点介绍、餐饮推荐和住宿建议。
-";
-    println!("\n用户: {prompt}");
-    print!("Assistant: ");
-    std::io::stdout().flush()?;
-
-    let config = GenerationConfig {
-        max_new_tokens: Some(1200),
-        temperature: 0.7,
-        top_p: 0.95,
-    };
-
-    let start = Instant::now();
-    let stream = model.generate_stream(prompt, false, &config)?;
-
-    for text in stream {
-        print!("{text}");
-        std::io::stdout().flush().ok();
-    }
-
-    let elapsed = start.elapsed();
-
-    println!();
-    println!("\n生成耗时: {:.2?}", elapsed);
-
-    Ok(())
-}

examples/wgpu-backend/Cargo.toml

@@ -5,5 +5,5 @@ edition = "2021"
 
 [dependencies]
 minicpm-inference = { path = "../../crates/minicpm-inference" }
-burn = { version = "0.21", default-features = false, features = ["std", "wgpu"] }
+burn = { version = "0.21", default-features = false, features = ["std", "wgpu", "fusion", "autotune"] }
 anyhow = "1.0"

examples/wgpu-backend/src/main.rs

@@ -1,10 +1,3 @@
-// Wgpu 后端推理示例 —— GPU 加速推理
-//
-// 依赖 WebGPU API，支持 CUDA/Metal/Vulkan 等平台。
-// 注意：与 Tauri 应用可能存在兼容性冲突，建议使用 Flex 后端。
-//
-// 运行方式：cargo run --release -p wgpu-backend [--full|--half|--q8]
-
 use burn::backend::{wgpu::WgpuDevice, Wgpu};
 use minicpm_inference::{GenerationConfig, MiniCPM};
 use std::io::Write;
@@ -12,49 +5,25 @@ use std::time::Instant;
 
 fn main() -> anyhow::Result<()> {
     let args: Vec<String> = std::env::args().collect();
-    let mode = args.get(1).map(|s| s.as_str()).unwrap_or("--full");
+    let think = args.iter().any(|s| s == "--think");
 
-    let (model_dir, config_path, tokenizer_path, is_q8) = match mode {
-        "--half" => (
-            "model_half",
-            "MiniCPM5-1B/config.json",
-            "MiniCPM5-1B/tokenizer.json",
-            false,
-        ),
-        "--q8" => (
-            "model_q8",
-            "MiniCPM5-1B/config.json",
-            "MiniCPM5-1B/tokenizer.json",
-            true,
-        ),
-        _ => (
-            "model",
-            "MiniCPM5-1B/config.json",
-            "MiniCPM5-1B/tokenizer.json",
-            false,
-        ),
-    };
-
-    println!("正在加载模型（后端：Wgpu，模式：{mode}）...");
-    println!("Wgpu 后端特性：GPU 加速（CUDA/Metal/Vulkan）");
-    let start = Instant::now();
-
-    let device = WgpuDevice::default();
-    let model: MiniCPM<Wgpu> = if is_q8 {
-        MiniCPM::<Wgpu>::load_q8(model_dir, config_path, tokenizer_path, &device)?
+    println!("========================================");
+    println!("后端: Wgpu (GPU 加速)");
+    println!("模型目录: model");
+    if think {
+        println!("思考模式: 开启");
     } else {
-        let model_path = format!("{model_dir}/model");
-        MiniCPM::<Wgpu>::load(&model_path, config_path, tokenizer_path, &device)?
-    };
+        println!("思考模式: 关闭");
+    }
+    println!("========================================");
 
+    let start = Instant::now();
+    let device = WgpuDevice::default();
+
+    let model = MiniCPM::<Wgpu>::load("model/model", "MiniCPM5-1B/config.json", "MiniCPM5-1B/tokenizer.json", &device)?;
     println!("模型加载完成，耗时: {:.2?}", start.elapsed());
 
-    let prompt = "需要处理生成一个河南的旅游计划，要求如下：
-1. 旅游计划为三天两夜的行程安排。
-2. 每天的行程安排包括景点、餐饮和住宿建议。
-3. 景点安排要考虑交通便利性和时间合理性。
-请生成详细的旅游计划，包含每天的行程安排、景点介绍、餐饮推荐和住宿建议。
-";
+    let prompt = "需要处理生成一个河南的旅游计划，要求如下：\n1. 旅游计划为三天两夜的行程安排。\n2. 每天的行程安排包括景点、餐饮和住宿建议。\n3. 景点安排要考虑交通便利性和时间合理性。\n请生成详细的旅游计划，包含每天的行程安排、景点介绍、餐饮推荐和住宿建议。\n";
     println!("\n用户: {prompt}");
     print!("Assistant: ");
     std::io::stdout().flush()?;
@@ -66,7 +35,7 @@ fn main() -> anyhow::Result<()> {
     };
 
     let start = Instant::now();
-    let stream = model.generate_stream(prompt, false, &config)?;
+    let stream = model.generate_stream(prompt, think, &config)?;
 
     for text in stream {
         print!("{text}");
@@ -74,7 +43,6 @@ fn main() -> anyhow::Result<()> {
     }
 
     let elapsed = start.elapsed();
-
     println!();
     println!("\n生成耗时: {:.2?}", elapsed);
 

model/config.json

@@ -0,0 +1,30 @@
+{
+    "_name_or_path": "openbmb/MiniCPM5-1B",
+    "architectures": [
+        "LlamaForCausalLM"
+    ],
+    "bos_token_id": 0,
+    "eos_token_id": [
+        1,
+        130073
+    ],
+    "pad_token_id": 1,
+    "hidden_act": "silu",
+    "hidden_size": 1536,
+    "initializer_range": 0.02,
+    "intermediate_size": 4608,
+    "max_position_embeddings": 131072,
+    "model_type": "llama",
+    "num_attention_heads": 16,
+    "num_hidden_layers": 24,
+    "num_key_value_heads": 2,
+    "head_dim": 128,
+    "rms_norm_eps": 1e-06,
+    "rope_theta": 5000000,
+    "rope_scaling": null,
+    "tie_word_embeddings": false,
+    "torch_dtype": "bfloat16",
+    "transformers_version": "5.6.2",
+    "use_cache": true,
+    "vocab_size": 130560
+}