uzu

一款面向 Apple Silicon 平台的高性能 AI 模型推理引擎。主要特性：

• 简单、高层次的 API
• 统一的模型配置，便于快速支持新模型
• 可追踪的计算过程，确保与权威实现的一致性
• 充分利用 Apple 设备上的统一内存

概述

有关架构的详细说明，请参阅 文档。

模型

uzu 使用自定义的模型格式。您可以使用示例脚本下载一个测试模型：

./scripts/download_test_model.sh

或者，您也可以使用以下命令下载任何已转换的支持模型：

cd ./tools/
uv run downloader list # 显示支持的模型列表
uv run downloader download {REPO_ID} # 使用 repo_id 下载特定模型

下载完成后，您可以在 ./workspace/models/{VERSION}/ 目录下找到该模型。

此外，您还可以使用 lalamo 自行导出特定模型：

git clone https://github.com/trymirai/lalamo.git
cd lalamo

之后，您可以查看支持的模型列表：

uv run lalamo list-models

然后导出所需的模型：

uv run lalamo convert meta-llama/Llama-3.2-1B-Instruct

绑定

• uzu-swift：预构建的 Swift 框架，可直接通过 SPM 使用。
• uzu-ts：为 Node.js 生态系统打造的预构建 TypeScript 框架。

命令行工具

您可以通过 CLI 模式运行 uzu：

cargo run --release -p cli -- help

用法：cli [COMMAND]
​
命令：
  run    运行指定路径的模型
  serve  启动指定模型路径的服务端
  bench  对指定模型进行基准测试
  help   打印此消息或给定子命令的帮助信息

编译

目前我们仅支持 Metal 后端，因此要编译相应的内核，您需要安装 Xcode 并执行以下命令：

xcodebuild -runFirstLaunch
xcodebuild -downloadComponent MetalToolchain

快速入门

首先，在您的 Cargo.toml 中添加 uzu 依赖：

[dependencies]
uzu = { git = "https://github.com/trymirai/uzu", branch = "main", package = "uzu" }

然后，使用特定的模型和配置创建一个推理会话：

use std::path::PathBuf;

use uzu::session::{
    Session,
    config::{DecodingConfig, RunConfig},
    types::{Input, Output},
};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    let model_path = PathBuf::from("MODEL_PATH");
    let mut session = Session::new(model_path, DecodingConfig::default())?;

    let input = Input::Text(String::from("Tell about London"));
    let output = session.run(
        input,
        RunConfig::default().tokens_limit(128),
        Some(|_: Output| {
            return true;
        }),
    )?;
    println!("{}", output.text.original);
    Ok(())
}

基准测试

要运行基准测试，可以使用以下命令：

cargo run --release -p cli -- bench ./workspace/models/{ENGINE_VERSION}/{MODEL_NAME} ./workspace/models/{ENGINE_VERSION}/{MODEL_NAME}/benchmark_task.json ./workspace/models/{ENGINE_VERSION}/{MODEL_NAME}/benchmark_result.json

benchmark_task.json 文件会在您通过前面提到的 ./tools/helpers/ 脚本下载模型后自动生成。

许可证

本项目采用 MIT 许可证。详情请参阅 LICENSE 文件。

uzu 快速上手指南

uzu 是一个专为 Apple Silicon 芯片设计的高性能 AI 模型推理引擎。它利用统一内存架构，提供简洁的 API 和可追踪的计算过程，适合在 macOS 设备上高效运行大语言模型。

环境准备

系统要求：

• 操作系统：macOS (必须运行在 Apple Silicon M1/M2/M3 等芯片上)
• 开发工具：Xcode (用于编译 Metal 后端内核)

前置依赖安装： 在开始之前，请确保已安装 Xcode 并配置好 Metal 工具链。执行以下命令：

xcodebuild -runFirstLaunch
xcodebuild -downloadComponent MetalToolchain

此外，项目使用 uv 管理 Python 工具链（用于模型下载），请确保已安装 uv。

安装步骤

1. 获取模型

uzu 使用专用的模型格式。你可以选择下载预转换的模型或自行转换。

方式 A：下载预支持模型（推荐） 进入工具目录并列出/下载模型：

cd ./tools/
uv run downloader list # 查看支持的模型列表
uv run downloader download {REPO_ID} # 替换 {REPO_ID} 下载特定模型

下载后的模型位于 ./workspace/models/{VERSION}/。

方式 B：自行转换模型 如果你需要转换其他模型，可以使用 lalamo 工具：

git clone https://github.com/trymirai/lalamo.git
cd lalamo
uv run lalamo list-models
uv run lalamo convert meta-llama/Llama-3.2-1B-Instruct

2. 集成到 Rust 项目

在你的 Rust 项目 Cargo.toml 中添加依赖：

[dependencies]
uzu = { git = "https://github.com/trymirai/uzu", branch = "main", package = "uzu" }

基本使用

以下是一个最简单的 Rust 代码示例，展示如何初始化会话并进行文本生成推理。

将 MODEL_PATH 替换为你实际下载的模型路径：

use std::path::PathBuf;

use uzu::session::{
    Session,
    config::{DecodingConfig, RunConfig},
    types::{Input, Output},
};

fn main() -> Result<(), Box<dyn std::error::Error>> {
    // 初始化会话
    let model_path = PathBuf::from("MODEL_PATH");
    let mut session = Session::new(model_path, DecodingConfig::default())?;

    // 准备输入
    let input = Input::Text(String::from("Tell about London"));
    
    // 运行推理
    let output = session.run(
        input,
        RunConfig::default().tokens_limit(128),
        Some(|_: Output| {
            return true;
        }),
    )?;
    
    // 输出结果
    println!("{}", output.text.original);
    Ok(())
}

运行项目即可看到模型生成的关于伦敦的介绍。

提示：你也可以直接使用 CLI 模式运行模型，无需编写代码：

cargo run --release -p cli -- run MODEL_PATH