jhqxxx
16 KiB
16 KiB
架构与设计
本文档深入介绍 AHA 的架构和设计原则。
概述
AHA（高性能 AI 推理引擎）是一个基于 Rust 的库，构建在 Candle 框架之上。它提供了统一的接口，可以在本地运行多个最先进的 AI 模型，无需 API 密钥或云服务。
核心特性
- 本地优先：所有推理都在您的机器上运行
- 多模态支持：支持文本、视觉、音频、OCR 和 ASR 模型
- 跨平台：支持 Linux、macOS 和 Windows
- GPU 加速：可选的 CUDA 和 Metal 支持
- 内存安全：使用 Rust 构建，确保安全性和性能
- OpenAI 兼容：易于与现有工具集成
架构原则
1. 本地优先设计
AHA 设计为完全在本地运行：
- 无云依赖：所有模型都在本地下载和运行
- 隐私保护：您的数据永远不会离开您的机器
- 无需 API 密钥：下载后模型可永久使用
- 离线可用：下载后模型无需互联网连接
2. 统一模型接口
所有模型都实现通用的 GenerateModel trait，提供：
- 不同模型类型之间的一致 API
- 无需更改代码即可轻松切换模型
- 支持实时输出的流式响应
- 标准化的错误处理
3. 跨平台支持
AHA 抽象了平台差异：
- 设备抽象：自动 CPU/GPU 检测和选择
- 精度处理：基于硬件动态选择 F32/F16/BF16
- 路径管理：跨平台一致的模型存储
核心组件
┌─────────────────────────────────────────────────────────────┐
│ CLI 层 │
│ (main.rs - 命令解析、模型下载、服务管理) │
└─────────────────────────────────────────────────────────────┘
│
┌─────────────────────────────────────────────────────────────┐
│ HTTP API 层 │
│ (api.rs - OpenAI 兼容端点、流式传输、认证) │
└─────────────────────────────────────────────────────────────┘
│
┌─────────────────────────────────────────────────────────────┐
│ 模型抽象层 │
│ (GenerateModel trait - 统一接口) │
└─────────────────────────────────────────────────────────────┘
│
┌─────────────────────┼─────────────────────┐
│ │ │
┌───────▼────────┐ ┌────────▼─────────┐ ┌───────▼────────┐
│ 文本模型 │ │ 视觉模型 │ │ 音频模型 │
│ - Qwen3 │ │ - Qwen2.5VL │ │ - VoxCPM │
│ - MiniCPM4 │ │ - Qwen3VL │ │ - VoxCPM1.5 │
└────────────────┘ └──────────────────┘ └────────────────┘
│ │ │
┌───────▼────────┐ ┌────────▼─────────┐ ┌───────▼────────┐
│ OCR 模型 │ │ ASR 模型 │ │ 图像模型 │
│ - DeepSeek │ │ - GLM-ASR │ │ - RMBG2.0 │
│ - Hunyuan │ │ - Fun-ASR │ │ │
│ - PaddleOCR │ │ - Qwen3-ASR │ │ │
└────────────────┘ └──────────────────┘ └────────────────┘
│
┌─────────────────────────────────────────────────────────────┐
│ 工具模块 │
│ - tokenizer: 分词工具 │
│ - chat_template: 聊天格式处理 │
│ - position_embed: 位置编码 │
│ - utils: 通用工具（音频、图像、下载） │
└─────────────────────────────────────────────────────────────┘
│
┌─────────────────────────────────────────────────────────────┐
│ Candle ML 框架 │
│ (张量操作、模型加载、设备管理) │
└─────────────────────────────────────────────────────────────┘
CLI 层 (src/main.rs)
CLI 层提供命令行界面功能：
- 命令解析：使用
clap进行参数解析 - 模型管理：自动下载和缓存
- 服务控制：启动/停止 HTTP 服务器
- 直接推理：无需服务器即可运行模型
可用命令：
cli- 下载模型并启动服务（默认）serv- 使用现有模型启动服务download- 仅下载模型run- 直接模型推理list- 列出支持的模型
HTTP API 层 (src/api.rs)
HTTP API 层提供 REST 端点：
- OpenAI 兼容：匹配 OpenAI API 格式
- 流式支持：实时响应生成
- 多模态：处理文本、图像和音频
- 线程安全：使用 RwLock 处理并发请求
端点：
POST /chat/completions- 聊天和文本生成POST /images/remove_background- 图像背景移除POST /audio/speech- 文本转语音合成
模型抽象层
所有模型都实现 GenerateModel trait：
pub trait GenerateModel {
// 生成响应
fn generate(&mut self, prompt: &str, params: GenerationParams) -> Result<String>;
// 流式生成
fn generate_stream(&mut self, prompt: &str, params: GenerationParams)
-> Result<Box<dyn Iterator<Item = Result<String>>>>;
}
这提供了：
- 多态性：统一处理不同模型
- 可扩展性：易于添加新模型
- 类型安全：编译时保证
工具模块
分词器 (src/tokenizer/)
- 从模型配置加载分词器
- 处理特殊标记
- 管理词汇表
聊天模板 (src/chat_template/)
- 将聊天消息格式化为模型提示
- 支持多种聊天格式（ChatML 等）
- 处理系统消息和角色标签
位置编码 (src/position_embed/)
- 为 transformer 实现位置编码
- 支持 RoPE（旋转位置编码）
- 处理多模态模型的 M-RoPE
工具 (src/utils/)
audio_utils.rs- 音频处理（WAV、MP3）image_utils.rs- 图像处理（调整大小、编码/解码）tensor_utils.rs- Tensor常用方法mod.rs- 通用工具和常量
设计模式
1. 基于 Trait 的抽象
GenerateModel trait 提供统一接口：
// 所有模型都实现此 trait
impl GenerateModel for Qwen3VL { /* ... */ }
impl GenerateModel for VoxCPM { /* ... */ }
impl GenerateModel for DeepSeekOCR { /* ... */ }
// 使用方式与模型无关
let mut model: Box<dyn GenerateModel> = load_model(model_type)?;
let result = model.generate(prompt, params)?;
2. 工厂模式
模型加载使用工厂函数：
pub fn load_model(
model_type: WhichModel,
model_path: &str,
device: &Device,
) -> Result<Box<dyn GenerateModel>> {
match model_type {
WhichModel::Qwen3VL2B => Ok(Box::new(qwen3vl::generate::Qwen3VLGenerate::init(...)?)),
WhichModel::VoxCPM1_5 => Ok(Box::new(voxcpm::generate::VoxCPMGenerate::init(...)?)),
// ... other models
_ => Err(anyhow!("Unsupported model: {}", model_type)),
}
}
3. 命令模式
CLI 子命令封装不同的操作：
match command {
Commands::Cli { model, port, address } => { /* 下载并启动服务 */ }
Commands::Serv { model, weight_path, port } => { /* 仅启动服务 */ }
Commands::Download { model, save_dir } => { /* 仅下载 */ }
Commands::Run { model, input, weight_path } => { /* 直接推理 */ }
Commands::List => { /* 列出模型 */ }
}
模型组织结构
每个模型都遵循一致的结构：
src/models/{model_name}/
├── config.rs # 模型配置和生成参数
├── model.rs # 核心模型架构（层、注意力）
├── generate.rs # 推理逻辑（实现 GenerateModel trait）
├── processor.rs # 模型特定处理（用于复杂模型）
└── mod.rs # 模块声明和导出
示例：Qwen3VL
src/models/qwen3vl/
├── config.rs # Qwen3VLConfig、GenerationConfig
├── model.rs # Qwen3VL transformer 层、注意力机制
├── generate.rs # Qwen3VLGenerate 实现
├── processor.rs # 多模态输入的图像和文本处理
└── mod.rs # 导出公共 API
性能优化
GPU 加速
AHA 通过以下方式支持 GPU 加速：
- CUDA：用于 NVIDIA GPU（Linux、Windows）
- Metal：用于 Apple Silicon（macOS）
启用方式：
cargo build --features cuda # NVIDIA GPU
cargo build --features metal # Apple Silicon
Flash Attention
Flash Attention 优化长序列处理：
- 减少内存使用
- 提高推理速度
- 对视觉模型特别有益
启用方式：
cargo build --features cuda,flash-attn
内存映射张量
模型使用内存映射文件：
- 更快的加载时间
- 减少内存占用
- 支持并发模型加载
精度优化
基于硬件的动态精度选择：
- F32：最高精度（仅 CPU）
- F16：平衡性能（GPU）
- BF16：最适合现代 GPU
安全考虑
本地处理
- 模型下载后无外部 API 调用
- 无遥测或数据收集
- 数据完全保留在本地
内存安全
- Rust 所有权系统防止内存泄漏
- 无缓冲区溢出或使用后释放错误
- 线程安全的并发操作
输入验证
- 文件大小限制（字符串 5MB，文件 100MB）
- 路径验证防止目录遍历
- 类型安全的请求处理
数据流
请求流程
┌─────────┐
│ 客户端 │
└────┬────┘
│ HTTP 请求
▼
┌──────────────────────────────────────────────────────────┐
│ Rocket HTTP 服务器 │
│ - 将请求路由到端点 │
│ - 解析请求体 │
│ - 提取参数 │
└────────────┬─────────────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────┐
│ API 处理器 (api.rs) │
│ - 获取模型锁 │
│ - 准备输入（分词、处理图像/音频） │
│ - 调用 model.generate() 或 generate_stream() │
└────────────┬─────────────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────┐
│ 模型实现 (models/{model}/generate.rs) │
│ - 从内存映射文件加载权重 │
│ - 通过 Candle 张量运行前向传播 │
│ - 解码输出标记 │
└────────────┬─────────────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────┐
│ Candle 框架 │
│ - 在 CPU 或 GPU 设备上执行 │
│ - 管理张量操作 │
└────────────┬─────────────────────────────────────────────┘
│
▼
┌──────────────────────────────────────────────────────────┐
│ 响应生成 │
│ - 格式化响应（JSON / 流式） │
│ - 返回给客户端 │
└──────────────────────────────────────────────────────────┘
模型加载流程
用户指定模型
│
▼
检查是否提供了 --weight-path
│
┌───┴───┐
│ │
是 否
│ │
▼ ▼
使用本地 从 ModelScope 下载
路径 │
│ ▼
│ 保存到 ~/.aha/{model}/
│ │
└───┬────┘
▼
将模型权重加载到内存
│
▼
初始化模型 (init())
│
▼
准备就绪，可以进行推理
扩展点
添加新模型
- 在
src/models/下创建模型目录 - 实现
GenerateModeltrait - 在
mod.rs的工厂函数中添加模型 - 在
main.rs中添加 CLI 映射 - 在
tests/中添加测试用例
自定义处理
模型可以覆盖默认处理：
- 自定义分词
- 特殊的输入/输出格式
- 模型特定的优化